-
Notifications
You must be signed in to change notification settings - Fork 661
/
uniform_resample.go
72 lines (61 loc) · 1.43 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
// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package sampler
import "golang.org/x/exp/maps"
// 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) {
s.rng = globalRNG
s.seededRNG = newRNG()
s.length = length
s.drawn = make(map[uint64]struct{})
}
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() {
maps.Clear(s.drawn)
}
func (s *uniformResample) Next() (uint64, error) {
i := uint64(len(s.drawn))
if i >= s.length {
return 0, ErrOutOfRange
}
for {
draw := s.rng.Uint64Inclusive(s.length - 1)
if _, ok := s.drawn[draw]; ok {
continue
}
s.drawn[draw] = struct{}{}
return draw, nil
}
}