-
Notifications
You must be signed in to change notification settings - Fork 658
/
weighted_linear.go
92 lines (74 loc) · 2.09 KB
/
weighted_linear.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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
// Copyright (C) 2019-2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package sampler
import (
"sort"
safemath "github.com/ava-labs/avalanchego/utils/math"
)
var _ Weighted = &weightedLinear{}
type weightedLinearElement struct {
cumulativeWeight uint64
index int
}
// Sampling is performed by executing a linear search over the provided elements
// in the order of their probabilistic occurrence.
//
// Initialization takes O(n * log(n)) time, where n is the number of elements
// that can be sampled.
// Sampling can take up to O(n) time. As the distribution becomes more biased,
// sampling will become faster in expectation.
type weightedLinear struct {
arr []weightedLinearElement
}
func (s *weightedLinear) Initialize(weights []uint64) error {
numWeights := len(weights)
if numWeights <= cap(s.arr) {
s.arr = s.arr[:numWeights]
} else {
s.arr = make([]weightedLinearElement, numWeights)
}
for i, weight := range weights {
s.arr[i] = weightedLinearElement{
cumulativeWeight: weight,
index: i,
}
}
// Optimize so that the most probable values are at the front of the array
sortWeightedLinear(s.arr)
for i := 1; i < len(s.arr); i++ {
newWeight, err := safemath.Add64(
s.arr[i-1].cumulativeWeight,
s.arr[i].cumulativeWeight,
)
if err != nil {
return err
}
s.arr[i].cumulativeWeight = newWeight
}
return nil
}
func (s *weightedLinear) Sample(value uint64) (int, error) {
if len(s.arr) == 0 || s.arr[len(s.arr)-1].cumulativeWeight <= value {
return 0, errOutOfRange
}
index := 0
for {
if elem := s.arr[index]; value < elem.cumulativeWeight {
return elem.index, nil
}
index++
}
}
type innerSortWeightedLinear []weightedLinearElement
func (lst innerSortWeightedLinear) Less(i, j int) bool {
return lst[i].cumulativeWeight > lst[j].cumulativeWeight
}
func (lst innerSortWeightedLinear) Len() int {
return len(lst)
}
func (lst innerSortWeightedLinear) Swap(i, j int) {
lst[j], lst[i] = lst[i], lst[j]
}
func sortWeightedLinear(lst []weightedLinearElement) {
sort.Sort(innerSortWeightedLinear(lst))
}