You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

107 lines
2.6 KiB

  1. package main
  2. import (
  3. "math/rand"
  4. "sort"
  5. )
  6. // combinations takes input in the form of a map of item lists, and returns a
  7. // list of all combinations of each item for each key. E.g.:
  8. //
  9. // {"foo": [1, 2, 3], "bar": [4, 5, 6]}
  10. //
  11. // Will return the following maps:
  12. //
  13. // {"foo": 1, "bar": 4}
  14. // {"foo": 1, "bar": 5}
  15. // {"foo": 1, "bar": 6}
  16. // {"foo": 2, "bar": 4}
  17. // {"foo": 2, "bar": 5}
  18. // {"foo": 2, "bar": 6}
  19. // {"foo": 3, "bar": 4}
  20. // {"foo": 3, "bar": 5}
  21. // {"foo": 3, "bar": 6}
  22. func combinations(items map[string][]interface{}) []map[string]interface{} {
  23. keys := []string{}
  24. for key := range items {
  25. keys = append(keys, key)
  26. }
  27. sort.Strings(keys)
  28. return combiner(map[string]interface{}{}, keys, items)
  29. }
  30. // combiner is a utility function for combinations.
  31. func combiner(head map[string]interface{}, pending []string, items map[string][]interface{}) []map[string]interface{} {
  32. if len(pending) == 0 {
  33. return []map[string]interface{}{head}
  34. }
  35. key, pending := pending[0], pending[1:]
  36. result := []map[string]interface{}{}
  37. for _, value := range items[key] {
  38. path := map[string]interface{}{}
  39. for k, v := range head {
  40. path[k] = v
  41. }
  42. path[key] = value
  43. result = append(result, combiner(path, pending, items)...)
  44. }
  45. return result
  46. }
  47. // uniformChoice chooses a single random item from the argument list, uniformly weighted.
  48. type uniformChoice []interface{}
  49. func (uc uniformChoice) Choose(r *rand.Rand) interface{} {
  50. return uc[r.Intn(len(uc))]
  51. }
  52. // weightedChoice chooses a single random key from a map of keys and weights.
  53. type weightedChoice map[interface{}]uint // nolint:unused
  54. func (wc weightedChoice) Choose(r *rand.Rand) interface{} {
  55. total := 0
  56. choices := make([]interface{}, 0, len(wc))
  57. for choice, weight := range wc {
  58. total += int(weight)
  59. choices = append(choices, choice)
  60. }
  61. rem := r.Intn(total)
  62. for _, choice := range choices {
  63. rem -= int(wc[choice])
  64. if rem <= 0 {
  65. return choice
  66. }
  67. }
  68. return nil
  69. }
  70. // probSetChoice picks a set of strings based on each string's probability (0-1).
  71. type probSetChoice map[string]float64
  72. func (pc probSetChoice) Choose(r *rand.Rand) []string {
  73. choices := []string{}
  74. for item, prob := range pc {
  75. if r.Float64() <= prob {
  76. choices = append(choices, item)
  77. }
  78. }
  79. return choices
  80. }
  81. // uniformSetChoice picks a set of strings with uniform probability, picking at least one.
  82. type uniformSetChoice []string
  83. func (usc uniformSetChoice) Choose(r *rand.Rand) []string {
  84. choices := []string{}
  85. indexes := r.Perm(len(usc))
  86. if len(indexes) > 1 {
  87. indexes = indexes[:1+r.Intn(len(indexes)-1)]
  88. }
  89. for _, i := range indexes {
  90. choices = append(choices, usc[i])
  91. }
  92. return choices
  93. }