combin: Add helpers for dealing with permutations (#1076)

* combin: Add helpers for dealing with permutations

stat/combin/combin.go

@@ -72,8 +72,8 @@ func LogGeneralizedBinomial(n, k float64) float64 {
 	return a - b - c
 }
 
-// CombinationGenerator generates combinations iteratively. Combinations may be
+// CombinationGenerator generates combinations iteratively. The Combinations
-// called to generate all combinations collectively.
+// function may be called to generate all combinations collectively.
 type CombinationGenerator struct {
 	n         int
 	k         int
@@ -118,29 +118,28 @@ func (c *CombinationGenerator) Next() bool {
 	return true
 }
 
-// Combination generates the current combination. If next is non-nil, it must have
+// Combination returns the current combination. If dst is non-nil, it must have
-// length k and the result will be stored in-place into combination. If combination
+// length k and the result will be stored in-place into dst. If dst
 // is nil a new slice will be allocated and returned. If all of the combinations
 // have already been constructed (Next() returns false), Combination will panic.
 //
 // Next must be called to initialize the first value before calling Combination
 // or Combination will panic. The value returned by Combination is only changed
 // during calls to Next.
-func (c *CombinationGenerator) Combination(combination []int) []int {
+func (c *CombinationGenerator) Combination(dst []int) []int {
 	if c.remaining == -1 {
 		panic("combin: all combinations have been generated")
 	}
 	if c.previous == nil {
 		panic("combin: Combination called before Next")
 	}
-	if combination == nil {
+	if dst == nil {
-		combination = make([]int, c.k)
+		dst = make([]int, c.k)
-	}
+	} else if len(dst) != c.k {
-	if len(combination) != c.k {
 		panic(badInput)
 	}
-	copy(combination, c.previous)
+	copy(dst, c.previous)
-	return combination
+	return dst
 }
 
 // Combinations generates all of the combinations of k elements from a
@@ -150,7 +149,7 @@ func (c *CombinationGenerator) Combination(combination []int) []int {
 // n and k must be non-negative with n >= k, otherwise Combinations will panic.
 //
 // CombinationGenerator may alternatively be used to generate the combinations
-// iteratively instead of collectively.
+// iteratively instead of collectively, or IndexToCombination for random access.
 func Combinations(n, k int) [][]int {
 	combins := Binomial(n, k)
 	data := make([][]int, combins)
@@ -368,3 +367,267 @@ func SubFor(sub []int, idx int, dims []int) []int {
 	sub[len(sub)-1] = idx
 	return sub
 }
+
+// NumPermutations returns the number of permutations when selecting k
+// objects from a set of n objects when the selection order matters.
+// No check is made for overflow.
+//
+// NumPermutations panics if either n or k is negative, or if k is
+// greater than n.
+func NumPermutations(n, k int) int {
+	if n < 0 {
+		panic("combin: n is negative")
+	}
+	if k < 0 {
+		panic("combin: k is negative")
+	}
+	if k > n {
+		panic("combin: k is greater than n")
+	}
+	p := 1
+	for i := n - k + 1; i <= n; i++ {
+		p *= i
+	}
+	return p
+}
+
+// Permutations generates all of the permutations of k elements from a
+// set of size n. The returned slice has length NumPermutations(n, k)
+// and each inner slice has length k.
+//
+// n and k must be non-negative with n >= k, otherwise Permutations will panic.
+//
+// PermutationGenerator may alternatively be used to generate the permutations
+// iteratively instead of collectively, or IndexToPermutation for random access.
+func Permutations(n, k int) [][]int {
+	nPerms := NumPermutations(n, k)
+	data := make([][]int, nPerms)
+	if len(data) == 0 {
+		return data
+	}
+	for i := 0; i < nPerms; i++ {
+		data[i] = IndexToPermutation(nil, i, n, k)
+	}
+	return data
+}
+
+// PermutationGenerator generates permutations iteratively. The Permutations
+// function may be called to generate all permutations collectively.
+type PermutationGenerator struct {
+	n           int
+	k           int
+	nPerm       int
+	idx         int
+	permutation []int
+}
+
+// NewPermutationGenerator returns a PermutationGenerator for generating the
+// permutations of k elements from a set of size n.
+//
+// n and k must be non-negative with n >= k, otherwise NewPermutationGenerator
+// will panic.
+func NewPermutationGenerator(n, k int) *PermutationGenerator {
+	return &PermutationGenerator{
+		n:           n,
+		k:           k,
+		nPerm:       NumPermutations(n, k),
+		idx:         -1,
+		permutation: make([]int, k),
+	}
+}
+
+// Next advances the iterator if there are permutations remaining to be generated,
+// and returns false if all permutations have been generated. Next must be called
+// to initialize the first value before calling Permutation or Permutation will
+// panic. The value returned by Permutation is only changed during calls to Next.
+func (p *PermutationGenerator) Next() bool {
+	if p.idx >= p.nPerm-1 {
+		p.idx = p.nPerm // so Permutation can panic.
+		return false
+	}
+	p.idx++
+	IndexToPermutation(p.permutation, p.idx, p.n, p.k)
+	return true
+}
+
+// Permutation returns the current permutation. If dst is non-nil, it must have
+// length k and the result will be stored in-place into dst. If dst
+// is nil a new slice will be allocated and returned. If all of the permutations
+// have already been constructed (Next() returns false), Permutation will panic.
+//
+// Next must be called to initialize the first value before calling Permutation
+// or Permutation will panic. The value returned by Permutation is only changed
+// during calls to Next.
+func (p *PermutationGenerator) Permutation(dst []int) []int {
+	if p.idx == p.nPerm {
+		panic("combin: all permutations have been generated")
+	}
+	if p.idx == -1 {
+		panic("combin: Permutation called before Next")
+	}
+	if dst == nil {
+		dst = make([]int, p.k)
+	} else if len(dst) != p.k {
+		panic(badInput)
+	}
+	copy(dst, p.permutation)
+	return dst
+}
+
+// PermutationIndex returns the index of the given permutation.
+//
+// The functions PermutationIndex and IndexToPermutation define a bijection
+// between the integers and the NumPermutations(n, k) number of possible permutations.
+// PermutationIndex returns the inverse of IndexToPermutation.
+//
+// PermutationIndex panics if perm is not a permutation of k of the first
+// [0,n) integers, if n or k are non-negative, or if k is greater than n.
+func PermutationIndex(perm []int, n, k int) int {
+	if n < 0 || k < 0 {
+		panic(badNegInput)
+	}
+	if n < k {
+		panic(badSetSize)
+	}
+	if len(perm) != k {
+		panic("combin: bad length permutation")
+	}
+	contains := make(map[int]struct{}, k)
+	for _, v := range perm {
+		if v < 0 || v >= n {
+			panic("combin: bad element")
+		}
+		contains[v] = struct{}{}
+	}
+	if len(contains) != k {
+		panic("combin: perm contains non-unique elements")
+	}
+	if n == k {
+		// The permutation is the ordering of the elements.
+		return equalPermutationIndex(perm)
+	}
+
+	// The permutation index is found by finding the combination index and the
+	// equalPermutation index. The combination index is found by just sorting
+	// the elements, and the the permutation index is the ordering of the size
+	// of the elements.
+	tmp := make([]int, len(perm))
+	copy(tmp, perm)
+	idx := make([]int, len(perm))
+	for i := range idx {
+		idx[i] = i
+	}
+	s := sortInts{tmp, idx}
+	sort.Sort(s)
+	order := make([]int, len(perm))
+	for i, v := range idx {
+		order[v] = i
+	}
+	combIdx := CombinationIndex(tmp, n, k)
+	permIdx := equalPermutationIndex(order)
+	return combIdx*NumPermutations(k, k) + permIdx
+}
+
+type sortInts struct {
+	data []int
+	idx  []int
+}
+
+func (s sortInts) Len() int {
+	return len(s.data)
+}
+
+func (s sortInts) Less(i, j int) bool {
+	return s.data[i] < s.data[j]
+}
+
+func (s sortInts) Swap(i, j int) {
+	s.data[i], s.data[j] = s.data[j], s.data[i]
+	s.idx[i], s.idx[j] = s.idx[j], s.idx[i]
+}
+
+// IndexToPermutation returns the permutation corresponding to the given index.
+//
+// The functions PermutationIndex and IndexToPermutation define a bijection
+// between the integers and the NumPermutations(n, k) number of possible permutations.
+// IndexToPermutation returns the inverse of PermutationIndex.
+//
+// The permutation is stored in-place into dst if dst is non-nil, otherwise
+// a new slice is allocated and returned.
+//
+// IndexToPermutation panics if n or k are non-negative, if k is greater than n,
+// or if idx is not in [0, NumPermutations(n,k)-1]. IndexToPermutation will also panic
+// if dst is non-nil and len(dst) is not k.
+func IndexToPermutation(dst []int, idx, n, k int) []int {
+	nPerm := NumPermutations(n, k)
+	if idx < 0 || idx >= nPerm {
+		panic("combin: invalid index")
+	}
+	if dst == nil {
+		dst = make([]int, k)
+	} else if len(dst) != k {
+		panic(badInput)
+	}
+	if n == k {
+		indexToEqualPermutation(dst, idx)
+		return dst
+	}
+
+	// First, we index into the combination (which of the k items to choose)
+	// and then we index into the n == k permutation of those k items. The
+	// indexing acts like a matrix with nComb rows and factorial(k) columns.
+	kPerm := NumPermutations(k, k)
+	combIdx := idx / kPerm
+	permIdx := idx % kPerm
+	comb := IndexToCombination(nil, combIdx, n, k) // Gives us the set of integers.
+	perm := make([]int, len(dst))
+	indexToEqualPermutation(perm, permIdx) // Gives their order.
+	for i, v := range perm {
+		dst[i] = comb[v]
+	}
+	return dst
+}
+
+// equalPermutationIndex returns the index of the given permutation of the
+// first k integers.
+func equalPermutationIndex(perm []int) int {
+	// Note(btracey): This is an n^2 algorithm, but factorial increases
+	// very quickly (25! overflows int64) so this is not a problem in
+	// practice.
+	idx := 0
+	for i, u := range perm {
+		less := 0
+		for _, v := range perm[i:] {
+			if v < u {
+				less++
+			}
+		}
+		idx += less * factorial(len(perm)-i-1)
+	}
+	return idx
+}
+
+// indexToEqualPermutation returns the permutation for the first len(dst)
+// integers for the given index.
+func indexToEqualPermutation(dst []int, idx int) {
+	for i := range dst {
+		dst[i] = i
+	}
+	for i := range dst {
+		f := factorial(len(dst) - i - 1)
+		r := idx / f
+		v := dst[i+r]
+		copy(dst[i+1:i+r+1], dst[i:i+r])
+		dst[i] = v
+		idx %= f
+	}
+}
+
+// factorial returns a!.
+func factorial(a int) int {
+	f := 1
+	for i := 2; i <= a; i++ {
+		f *= i
+	}
+	return f
+}

stat/combin/combin_test.go

@@ -297,3 +297,45 @@ func TestCartesian(t *testing.T) {
 		t.Errorf("cartesian data mismatch.\nwant:\n%v\ngot:\n%v", want, got)
 	}
 }
+
+func TestPermutationIndex(t *testing.T) {
+	for cas, s := range []struct {
+		n, k int
+	}{
+		{6, 3},
+		{4, 4},
+		{10, 1},
+		{8, 2},
+	} {
+		n := s.n
+		k := s.k
+		perms := make(map[string]struct{})
+		for i := 0; i < NumPermutations(n, k); i++ {
+			perm := IndexToPermutation(nil, i, n, k)
+			idx := PermutationIndex(perm, n, k)
+			if idx != i {
+				t.Errorf("Cas %d: permutation mismatch. Want %d, got %d", cas, i, idx)
+			}
+			perms[intSliceToKey(perm)] = struct{}{}
+		}
+		if len(perms) != NumPermutations(n, k) {
+			t.Errorf("Case %d: not all generated combinations were unique", cas)
+		}
+	}
+}
+
+func TestPermutationGenerator(t *testing.T) {
+	for n := 0; n <= 7; n++ {
+		for k := 1; k <= n; k++ {
+			permutations := Permutations(n, k)
+			pg := NewPermutationGenerator(n, k)
+			genPerms := make([][]int, 0, len(permutations))
+			for pg.Next() {
+				genPerms = append(genPerms, pg.Permutation(nil))
+			}
+			if !intSosMatch(permutations, genPerms) {
+				t.Errorf("Permutations and generated permutations do not match. n = %v, k = %v", n, k)
+			}
+		}
+	}
+}

stat/combin/combinations_example_test.go

@@ -10,26 +10,6 @@ import (
 	"gonum.org/v1/gonum/stat/combin"
 )
 
-func ExampleCombinations_Index() {
-	data := []string{"a", "b", "c", "d", "e"}
-	cs := combin.Combinations(len(data), 2)
-	for _, c := range cs {
-		fmt.Printf("%s%s\n", data[c[0]], data[c[1]])
-	}
-
-	// Output:
-	// ab
-	// ac
-	// ad
-	// ae
-	// bc
-	// bd
-	// be
-	// cd
-	// ce
-	// de
-}
-
 func ExampleCombinations() {
 	// combin provides several ways to work with the combinations of
 	// different objects. Combinations generates them directly.
@@ -40,13 +20,6 @@ func ExampleCombinations() {
 	for i, v := range list {
 		fmt.Println(i, v)
 	}
-	// The returned values can be used to index into a data structure.
-	data := []string{"a", "b", "c", "d", "e"}
-	cs := combin.Combinations(len(data), 2)
-	fmt.Println("\nString combinations:")
-	for _, c := range cs {
-		fmt.Printf("%s%s\n", data[c[0]], data[c[1]])
-	}
 	// This is easy, but the number of combinations  can be very large,
 	// and generating all at once can use a lot of memory.
 
@@ -62,8 +35,18 @@ func ExampleCombinations() {
 	// 7 [1 2 4]
 	// 8 [1 3 4]
 	// 9 [2 3 4]
-	//
+}
-	// String combinations:
+
+func ExampleCombinations_index() {
+	// The integer slices returned from Combinations can be used to index
+	// into a data structure.
+	data := []string{"a", "b", "c", "d", "e"}
+	cs := combin.Combinations(len(data), 2)
+	for _, c := range cs {
+		fmt.Printf("%s%s\n", data[c[0]], data[c[1]])
+	}
+
+	// Output:
 	// ab
 	// ac
 	// ad
@@ -74,7 +57,6 @@ func ExampleCombinations() {
 	// cd
 	// ce
 	// de
-
 }
 
 func ExampleCombinationGenerator() {
@@ -128,3 +110,149 @@ func ExampleIndexToCombination() {
 	// 8 [1 3 4] 8
 	// 9 [2 3 4] 9
 }
+
+func ExamplePermutations() {
+	// combin provides several ways to work with the permutationss of
+	// different objects. Permutations generates them directly.
+	fmt.Println("Generate list:")
+	n := 4
+	k := 3
+	list := combin.Permutations(n, k)
+	for i, v := range list {
+		fmt.Println(i, v)
+	}
+	// This is easy, but the number of permutations can be very large,
+	// and generating all at once can use a lot of memory.
+
+	// Output:
+	// Generate list:
+	// 0 [0 1 2]
+	// 1 [0 2 1]
+	// 2 [1 0 2]
+	// 3 [1 2 0]
+	// 4 [2 0 1]
+	// 5 [2 1 0]
+	// 6 [0 1 3]
+	// 7 [0 3 1]
+	// 8 [1 0 3]
+	// 9 [1 3 0]
+	// 10 [3 0 1]
+	// 11 [3 1 0]
+	// 12 [0 2 3]
+	// 13 [0 3 2]
+	// 14 [2 0 3]
+	// 15 [2 3 0]
+	// 16 [3 0 2]
+	// 17 [3 2 0]
+	// 18 [1 2 3]
+	// 19 [1 3 2]
+	// 20 [2 1 3]
+	// 21 [2 3 1]
+	// 22 [3 1 2]
+	// 23 [3 2 1]
+}
+
+func ExamplePermutations_index() {
+	// The integer slices returned from Permutations can be used to index
+	// into a data structure.
+	data := []string{"a", "b", "c", "d"}
+	cs := combin.Permutations(len(data), 2)
+	for _, c := range cs {
+		fmt.Printf("%s%s\n", data[c[0]], data[c[1]])
+	}
+
+	// Output:
+	// ab
+	// ba
+	// ac
+	// ca
+	// ad
+	// da
+	// bc
+	// cb
+	// bd
+	// db
+	// cd
+	// dc
+}
+
+func ExamplePermutationGenerator() {
+	// combin provides several ways to work with the permutations of
+	// different objects. PermutationGenerator constructs an iterator
+	// for the permutations.
+	n := 4
+	k := 3
+	gen := combin.NewPermutationGenerator(n, k)
+	idx := 0
+	for gen.Next() {
+		fmt.Println(idx, gen.Permutation(nil)) // can also store in-place.
+		idx++
+	}
+
+	// Output:
+	// 0 [0 1 2]
+	// 1 [0 2 1]
+	// 2 [1 0 2]
+	// 3 [1 2 0]
+	// 4 [2 0 1]
+	// 5 [2 1 0]
+	// 6 [0 1 3]
+	// 7 [0 3 1]
+	// 8 [1 0 3]
+	// 9 [1 3 0]
+	// 10 [3 0 1]
+	// 11 [3 1 0]
+	// 12 [0 2 3]
+	// 13 [0 3 2]
+	// 14 [2 0 3]
+	// 15 [2 3 0]
+	// 16 [3 0 2]
+	// 17 [3 2 0]
+	// 18 [1 2 3]
+	// 19 [1 3 2]
+	// 20 [2 1 3]
+	// 21 [2 3 1]
+	// 22 [3 1 2]
+	// 23 [3 2 1]
+}
+
+func ExampleIndexToPermutation() {
+	// combin provides several ways to work with the permutations of
+	// different objects. IndexToPermutation allows random access into
+	// the permutation order. Combined with PermutationIndex this
+	// provides a correspondence between integers and permutations.
+	n := 4
+	k := 3
+	comb := make([]int, k)
+	for i := 0; i < combin.NumPermutations(n, k); i++ {
+		combin.IndexToPermutation(comb, i, n, k) // can also use nil.
+		idx := combin.PermutationIndex(comb, n, k)
+		fmt.Println(i, comb, idx)
+	}
+
+	// Output:
+	// 0 [0 1 2] 0
+	// 1 [0 2 1] 1
+	// 2 [1 0 2] 2
+	// 3 [1 2 0] 3
+	// 4 [2 0 1] 4
+	// 5 [2 1 0] 5
+	// 6 [0 1 3] 6
+	// 7 [0 3 1] 7
+	// 8 [1 0 3] 8
+	// 9 [1 3 0] 9
+	// 10 [3 0 1] 10
+	// 11 [3 1 0] 11
+	// 12 [0 2 3] 12
+	// 13 [0 3 2] 13
+	// 14 [2 0 3] 14
+	// 15 [2 3 0] 15
+	// 16 [3 0 2] 16
+	// 17 [3 2 0] 17
+	// 18 [1 2 3] 18
+	// 19 [1 3 2] 19
+	// 20 [2 1 3] 20
+	// 21 [2 3 1] 21
+	// 22 [3 1 2] 22
+	// 23 [3 2 1] 23
+}