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.

85 lines
2.1 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. // probSetChoice picks a set of strings based on each string's probability (0-1).
  53. type probSetChoice map[string]float64
  54. func (pc probSetChoice) Choose(r *rand.Rand) []string {
  55. choices := []string{}
  56. for item, prob := range pc {
  57. if r.Float64() <= prob {
  58. choices = append(choices, item)
  59. }
  60. }
  61. return choices
  62. }
  63. // uniformSetChoice picks a set of strings with uniform probability, picking at least one.
  64. type uniformSetChoice []string
  65. func (usc uniformSetChoice) Choose(r *rand.Rand) []string {
  66. choices := []string{}
  67. indexes := r.Perm(len(usc))
  68. if len(indexes) > 1 {
  69. indexes = indexes[:1+r.Intn(len(indexes)-1)]
  70. }
  71. for _, i := range indexes {
  72. choices = append(choices, usc[i])
  73. }
  74. return choices
  75. }