-
Notifications
You must be signed in to change notification settings - Fork 0
/
comb.go
136 lines (117 loc) · 2.87 KB
/
comb.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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
package comb
var maxSizes = []uint64{0, ^uint64(0), 4294967296, 33290221, 102570, 13467, 3612, 1449, 746, 453, 308, 227, 178, 147, 125, 110, 99, 90, 84, 80, 75, 72, 69, 68, 66, 65, 64, 63, 63, 62, 62, 62}
const largestK = 31
const maxInt = uint64(^uint(0) >> 1)
func addHasOverflowed(a, b int) (sum int, overflow bool) {
sum = a + b
if (sum^a)&(sum^b) < 0 {
return sum, true
}
return sum, false
}
//CoeffUint64 returns the binomial coefficient n choose k as a uint64.
//CoeffUint64 panics if the calculation would overflow the uint64.
func CoeffUint64(n, k uint64) uint64 {
if k > n {
return 0
}
if k > n/2 {
k = n - k
}
if k == 0 {
return 1
}
if k > largestK || n > maxSizes[k] {
panic("calculation overflows uint64")
}
var comb uint64 = 1
var i uint64
for i = 1; i <= k; i++ {
comb *= (n - k + i)
comb /= i
}
return comb
}
//Coeff calculates the binomial coefficient n choose k and returns it as an int.
//Coeff panics if the calculation would overflow.
func Coeff(n, k int) int {
if n < 0 {
panic("n must be non-negative")
}
if k < 0 {
return 0
}
comb := CoeffUint64(uint64(n), uint64(k))
if comb > maxInt {
panic("coeff does not fit in an int")
}
return int(comb)
}
//Ranking and Unranking n choose k.
//Rank converts a sorted set of distinct ints to an int such that the set of ints can be recovered by Unrank.
//Rank panics if calculating the rank would cause an overflow.
func Rank(comb []int) int {
rank := 0
var overflow bool
for i, v := range comb {
c := Coeff(v, i+1)
rank, overflow = addHasOverflowed(rank, c)
if overflow {
panic("rank has overflowed int")
}
}
return rank
}
//Unrank converts an integer rank to a sorted set of k distinct ints. It is the inverse of rank.
func Unrank(rank, k int) []int {
comb := make([]int, k)
m := rank
for i := k - 1; i >= 0; i-- {
l := i + 1
b := 1
for b <= m {
b *= (l + 1)
b /= (l - i)
l++
}
comb[i] = l - 1
b *= (l - 1 - i)
b /= l
m -= b
}
return comb
}
//Iterator iterates over all ways of choosing k elements from 0, ..., n-1 in lexicographic order.
type Iterator struct {
n int
k int
data []int
}
//Next moves the Iterator to the next combination, returning true if there is one and false is there isn't.
func (b *Iterator) Next() bool {
for i := b.k - 1; i >= 0; i-- {
if b.data[i] < b.n+i-b.k {
b.data[i]++
for j := i + 1; j < b.k; j++ {
b.data[j] = b.data[j-1] + 1
}
return true
}
}
return false
}
//Value returns the current state of the iterator. This must not be modified.
func (b Iterator) Value() []int {
return b.data
}
//NewIterator returns a new Iterator to iterate over all subsets of k distinct elements from 0, ..., n-1 in lexicographic order.
func NewIterator(n, k int) *Iterator {
data := make([]int, k)
for i := 0; i < k; i++ {
data[i] = i
}
if k > 0 {
data[k-1]--
}
return &Iterator{n: n, k: k, data: data}
}