forked from dropbox/godropbox
/
rand.go
79 lines (65 loc) · 1.53 KB
/
rand.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
// rand2 is a collection of functions meant to supplement the capabilities
// provided by the standard "math/rand" package.
package rand2
import (
"math/rand"
"sort"
"github.com/dropbox/godropbox/container/set"
"github.com/dropbox/godropbox/errors"
)
// Samples 'k' unique ints from the range [0, n)
func SampleInts(n int, k int) (res []int, err error) {
if k < 0 {
err = errors.Newf("invalid sample size k")
return
}
if n < k {
err = errors.Newf("sample size k larger than n")
return
}
picked := set.NewSet()
for picked.Len() < k {
i := rand.Intn(n)
picked.Add(i)
}
res = make([]int, k)
e := 0
for i := range picked.Iter() {
res[e] = i.(int)
e++
}
return
}
// Samples 'k' elements from the given slice
func Sample(population []interface{}, k int) (res []interface{}, err error) {
n := len(population)
idxs, err := SampleInts(n, k)
if err != nil {
return
}
res = []interface{}{}
for _, idx := range idxs {
res = append(res, population[idx])
}
return
}
// Same as 'Sample' except it returns both the 'picked' sample set and the 'remaining' elements.
func PickN(population []interface{}, n int) (
picked []interface{}, remaining []interface{}, err error) {
total := len(population)
idxs, err := SampleInts(total, n)
if err != nil {
return
}
sort.Ints(idxs)
picked, remaining = []interface{}{}, []interface{}{}
for x, elem := range population {
if len(idxs) > 0 && x == idxs[0] {
picked = append(picked, elem)
idxs = idxs[1:]
} else {
remaining = append(remaining, elem)
}
}
return
}