-
Notifications
You must be signed in to change notification settings - Fork 84
/
stwHeap.go
61 lines (54 loc) · 1.3 KB
/
stwHeap.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
package alg
import (
"math/rand"
"github.com/oakmound/oak/dlog"
)
type stwHeap struct {
bh []float64
weightsBelow []float64
}
// Select Total Weight Heap
// This name was chosen relatively arbitrarily, if there
// is a canonical academic name for this structure we'd gladly
// use that instead
func newSTWHeap(f []float64) *stwHeap {
stwh := new(stwHeap)
f = append([]float64{0}, f...)
// The order of elements literally does not
// matter, so 'heap' is a misnomer.
stwh.bh = f
stwh.weightsBelow = make([]float64, len(f))
copy(stwh.weightsBelow, f)
for i := len(f) - 1; i > 1; i-- {
stwh.weightsBelow[i>>1] += stwh.weightsBelow[i]
}
return stwh
}
func (stwh *stwHeap) Pop() int {
if stwh.weightsBelow[1] <= ε {
dlog.Warn("Pop on stwHeap with no remaining elements")
return -1
}
w := stwh.weightsBelow[1] * rand.Float64()
i := 1
// With the >= here, we don't accept 0 weights
for w >= stwh.bh[i] {
w -= stwh.bh[i]
i <<= 1 // Propagate to left child
if w >= stwh.weightsBelow[i] {
w -= stwh.weightsBelow[i]
i++ // Switch to right child
}
}
i2 := i
w = stwh.bh[i]
// Instead of removing a node we set its weight to 0.
stwh.bh[i] = 0
// All parents of the index need to be reduced
// in total weight.
for i > 0 {
stwh.weightsBelow[i] -= w
i >>= 1
}
return i2 - 1
}