-
Notifications
You must be signed in to change notification settings - Fork 2
/
WeightedRandomChoice.go
73 lines (66 loc) · 1.69 KB
/
WeightedRandomChoice.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
package WeightedRandomChoice
import (
"math/rand"
"sort"
"time"
)
type WeightedRandomChoice struct {
Elements [] string
Weights [] int
TotalWeight int
calibratedWeights bool
precision int
calibrateValue int
}
func (wrc *WeightedRandomChoice) AddElement(element string, weight int) {
weight *= wrc.calibrateValue
i := sort.Search(len(wrc.Weights), func(i int) bool { return wrc.Weights[i] > weight })
wrc.Weights = append(wrc.Weights, 0)
wrc.Elements = append(wrc.Elements, "")
copy(wrc.Weights[i+1:], wrc.Weights[i:])
copy(wrc.Elements[i+1:], wrc.Elements[i:])
wrc.Weights[i] = weight
wrc.Elements[i] = element
wrc.TotalWeight += weight
}
func (wrc *WeightedRandomChoice) AddElements(elements map[string]int) {
for element, weight := range elements{
wrc.AddElement(element, weight)
}
}
func (wrc *WeightedRandomChoice) GetRandomChoice() string {
rand.Seed(time.Now().UnixNano())
value := rand.Intn(wrc.TotalWeight)
if !wrc.calibratedWeights {
wrc.calibrateWeights()
}
for key, weight := range wrc.Weights {
value -= weight
if value <= 0 {
return wrc.Elements[key]
}
}
return ""
}
func (wrc *WeightedRandomChoice) calibrateWeights() {
if wrc.TotalWeight/wrc.precision < 1 {
wrc.calibrateValue = wrc.precision / wrc.TotalWeight
wrc.TotalWeight = 0
for key := range wrc.Weights {
wrc.Weights[key] *= wrc.calibrateValue
wrc.TotalWeight += wrc.Weights[key]
}
wrc.calibratedWeights = true
}
}
func New(arguments ...int) WeightedRandomChoice {
var precision = 1000
if len(arguments) > 0 {
precision = arguments[0]
}
return WeightedRandomChoice{
precision: precision,
calibratedWeights: false,
calibrateValue: 1,
}
}