-
Notifications
You must be signed in to change notification settings - Fork 0
/
sample.go
90 lines (74 loc) · 1.74 KB
/
sample.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
package sampler
import (
"sync"
"golang.org/x/exp/slices"
)
// Sample returns a random selection of n items from items.
// Items with unique indexes are selected.
//
// [1 2 3 4 5] -> always unique items in result
// [1 2 3 4 4] -> 4 can be selected twice
func Sample[E any](rnd Rand, items []E, n int) []E {
return SampleAppend(rnd, nil, items, n)
}
// SampleAppend returns a random selection of n items from items.
// Items are appended to dst, which is grown if necessary.
// Items with unique indexes are selected.
//
// [1 2 3 4 5] -> always unique items in result
// [1 2 3 4 4] -> 4 can be selected twice
func SampleAppend[E any](rnd Rand, dst, items []E, n int) []E {
if n == 0 {
return dst
}
if n < 0 || n >= len(items) {
return fullShuffle(rnd, dst, items)
}
if dst == nil {
dst = make([]E, 0, n)
}
if 3*n > len(items) {
return permutate(rnd, dst, items, n)
}
set := make(map[int]struct{}, n)
for len(set) < n {
i := rnd.IntN(len(items))
if _, ok := set[i]; !ok {
set[i] = struct{}{}
dst = append(dst, items[i])
}
}
return dst
}
func permutate[E any](rnd Rand, dst, items []E, n int) []E {
ii := getIndexes(n)
defer releaseIndexes(ii)
perm[int](rnd, ii.values)
for _, i := range ii.values {
dst = append(dst, items[i])
}
return dst
}
type indexes struct {
values []int
}
var indexesPool = &sync.Pool{
New: func() any {
return &indexes{
values: make([]int, 64),
}
},
}
func getIndexes(n int) *indexes {
//nolint:errcheck // ignore casting error - we know that we have *indexes
ii := indexesPool.Get().(*indexes)
ii.values = slices.Grow(ii.values, n)[:n]
return ii
}
func releaseIndexes(ii *indexes) {
if len(ii.values) > 1024*1024 {
return
}
ii.values = ii.values[:0]
indexesPool.Put(ii)
}