-
Notifications
You must be signed in to change notification settings - Fork 658
/
uniform_resample.go
80 lines (69 loc) · 1.52 KB
/
uniform_resample.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
// Copyright (C) 2019-2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package sampler
import (
"math"
)
// uniformResample allows for sampling over a uniform distribution without
// replacement.
//
// Sampling is performed by sampling with replacement and resampling if a
// duplicate is sampled.
//
// Initialization takes O(1) time.
//
// Sampling is performed in O(count) time and O(count) space.
type uniformResample struct {
rng rng
seededRNG rng
length uint64
drawn map[uint64]struct{}
}
func (s *uniformResample) Initialize(length uint64) error {
if length > math.MaxInt64 {
return errOutOfRange
}
s.rng = globalRNG
s.seededRNG = newRNG()
s.length = length
s.drawn = make(map[uint64]struct{})
return nil
}
func (s *uniformResample) Sample(count int) ([]uint64, error) {
s.Reset()
results := make([]uint64, count)
for i := 0; i < count; i++ {
ret, err := s.Next()
if err != nil {
return nil, err
}
results[i] = ret
}
return results, nil
}
func (s *uniformResample) Seed(seed int64) {
s.rng = s.seededRNG
s.rng.Seed(seed)
}
func (s *uniformResample) ClearSeed() {
s.rng = globalRNG
}
func (s *uniformResample) Reset() {
for k := range s.drawn {
delete(s.drawn, k)
}
}
func (s *uniformResample) Next() (uint64, error) {
i := uint64(len(s.drawn))
if i >= s.length {
return 0, errOutOfRange
}
for {
draw := uint64(s.rng.Int63n(int64(s.length)))
if _, ok := s.drawn[draw]; ok {
continue
}
s.drawn[draw] = struct{}{}
return draw, nil
}
}