-
Notifications
You must be signed in to change notification settings - Fork 0
/
partitions.go
99 lines (91 loc) · 2.28 KB
/
partitions.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
package comb
func totalLength(a [][]int) int {
n := 0
for _, s := range a {
n += len(s)
}
return n
}
//The restricted growth string of a partition is the sequence of numbers r[0], ..., r[n-1] where r[i] is the index of the set that i is in when the sets are ordered in increasing size of their smallest element.
//sets must be ordered in increasing size of their smallest element.
func restrictedGrowthString(sets [][]int) []int {
rgs := make([]int, totalLength(sets))
curr := 0
for _, set := range sets {
for _, j := range set {
rgs[j] = curr
}
curr++
}
return rgs
}
func partitionFromRestrictedGrowthString(rgs []int) [][]int {
n := len(rgs)
max := 0
sizes := make([]int, n)
for _, v := range rgs {
sizes[v]++
if v > max {
max = v
}
}
p := make([][]int, max+1)
for i := range p {
p[i] = make([]int, 0, sizes[i])
}
for i, v := range rgs {
p[v] = append(p[v], i)
}
return p
}
//PartitionIterator iterates over all partitions of the set {0, ..., n-1}.
//The partitions are generated in lexicographic order of their restricted growth strings. It is safe to modify the output of .Value().
type PartitionIterator struct {
n int
m int
a []int
b []int
}
//TODO Handle n = 0
//NewPartitionIterator returns a *PartitionIterator which iterates over all partitions of the set {0, ..., n-1}.
//The partitions are generated in lexicographic order of their restricted growth sequences.
func NewPartitionIterator(n int) *PartitionIterator {
if n < 1 {
panic("Cannot handle n < 1")
}
a := make([]int, n)
a[n-1] = -1
b := make([]int, n)
for i := range b {
b[i] = 1
}
return &PartitionIterator{n: n, m: 1, a: a, b: b}
}
//Next tries to advance pi to the next partition, returning true if there is one and false if there isn't.
func (pi *PartitionIterator) Next() bool {
if pi.a[pi.n-1] == pi.m {
for j := pi.n - 2; j >= 1; j-- {
if pi.a[j] != pi.b[j] {
pi.a[j]++
pi.m = pi.b[j]
if pi.a[j] == pi.b[j] {
pi.m++
}
for k := j + 1; k < pi.n-1; k++ {
pi.a[k] = 0
pi.b[k] = pi.m
}
pi.a[pi.n-1] = 0
return true
}
}
return false
}
pi.a[pi.n-1]++
return true
}
//Value returns the current partition.
//It is safe to modify the output.
func (pi PartitionIterator) Value() [][]int {
return partitionFromRestrictedGrowthString(pi.a)
}