stat/combin: Add CombinationToIndex and IndexToCombination functions (#1054)

* stat/combin: Add CombinationToIndex and IndexToCombination functions

stat/combin/combin.go

@@ -6,6 +6,7 @@ package combin
 
 import (
 	"math"
+	"sort"
 )
 
 const (
@@ -183,7 +184,95 @@ func nextCombination(s []int, n, k int) {
 	}
 }
 
-// Cartesian returns indices into the cartesian product for sets of the given lengths. The Cartesian
+// CombinationIndex returns the index of the given combination.
+//
+// The functions CombinationIndex and IndexToCombination define a bijection
+// between the integers and the Binomial(n, k) number of possible combinations.
+// CombinationIndex returns the inverse of IndexToCombination.
+//
+// CombinationIndex panics if comb is not a sorted combination of the first
+// [0,n) integers, if n or k are non-negative, or if k is greater than n.
+func CombinationIndex(comb []int, n, k int) int {
+	if n < 0 || k < 0 {
+		panic(badNegInput)
+	}
+	if n < k {
+		panic(badSetSize)
+	}
+	if len(comb) != k {
+		panic("combin: bad length combination")
+	}
+	if !sort.IntsAreSorted(comb) {
+		panic("combin: input combination is not sorted")
+	}
+	contains := make(map[int]struct{}, k)
+	for _, v := range comb {
+		contains[v] = struct{}{}
+	}
+	if len(contains) != k {
+		panic("combin: comb contains non-unique elements")
+	}
+	// This algorithm iterates in reverse lexicograhpic order.
+	// Flip the index and values to swap the order.
+	rev := make([]int, k)
+	for i, v := range comb {
+		rev[len(comb)-i-1] = n - v - 1
+	}
+	idx := 0
+	for i, v := range rev {
+		if v >= i+1 {
+			idx += Binomial(v, i+1)
+		}
+	}
+	return Binomial(n, k) - 1 - idx
+}
+
+// IndexToCombination returns the combination corresponding to the given index.
+//
+// The functions CombinationIndex and IndexToCombination define a bijection
+// between the integers and the Binomial(n, k) number of possible combinations.
+// IndexToCombination returns the inverse of CombinationIndex (up to the order
+// of the elements).
+//
+// The combination is stored in-place into dst if dst is non-nil, otherwise
+// a new slice is allocated and returned.
+//
+// IndexToCombination panics if n or k are non-negative, if k is greater than n,
+// or if idx is not in [0, Binomial(n,k)-1]. IndexToCombination will also panic
+// if dst is non-nil and len(dst) is not k.
+func IndexToCombination(dst []int, idx, n, k int) []int {
+	if idx < 0 || idx >= Binomial(n, k) {
+		panic("combin: invalid index")
+	}
+	if dst == nil {
+		dst = make([]int, k)
+	} else if len(dst) != k {
+		panic(badInput)
+	}
+	// The base algorithm indexes in reverse lexicographic order
+	// flip the values and the index.
+	idx = Binomial(n, k) - 1 - idx
+	for i := range dst {
+		// Find the largest number m such that Binomial(m, k-i) <= idx.
+		// This is one less than the first number such that it is larger.
+		m := sort.Search(n, func(m int) bool {
+			if m < k-i {
+				return false
+			}
+			return Binomial(m, k-i) > idx
+		})
+		m--
+		// Normally this is put m into the last free spot, but we
+		// reverse the index and the value.
+		dst[i] = n - m - 1
+		if m >= k-i {
+			idx -= Binomial(m, k-i)
+		}
+	}
+	return dst
+}
+
+// Cartesian returns the cartesian product of the slices in data. The Cartesian
 // product of two sets is the set of all combinations of the items. For example,
 // given the input
 //  []int{2, 3, 1}

stat/combin/combin_example_test.go

@@ -1,31 +0,0 @@
-// Copyright ©2019 The Gonum Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file.
-
-package combin_test
-
-import (
-	"fmt"
-
-	"gonum.org/v1/gonum/stat/combin"
-)
-
-func ExampleCombinations() {
-	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
-}

stat/combin/combin_test.go

@@ -7,6 +7,7 @@ package combin
 import (
 	"math/big"
 	"reflect"
+	"strconv"
 	"testing"
 
 	"gonum.org/v1/gonum/floats"
@@ -181,38 +182,52 @@ func TestCombinationGenerator(t *testing.T) {
 	}
 }
 
-func TestCartesian(t *testing.T) {
-	// First, test with a known return.
-	lens := []int{2, 3, 4}
-	want := [][]int{
-		{0, 0, 0},
-		{0, 0, 1},
-		{0, 0, 2},
-		{0, 0, 3},
-		{0, 1, 0},
-		{0, 1, 1},
-		{0, 1, 2},
-		{0, 1, 3},
-		{0, 2, 0},
-		{0, 2, 1},
-		{0, 2, 2},
-		{0, 2, 3},
-		{1, 0, 0},
-		{1, 0, 1},
-		{1, 0, 2},
-		{1, 0, 3},
-		{1, 1, 0},
-		{1, 1, 1},
-		{1, 1, 2},
-		{1, 1, 3},
-		{1, 2, 0},
-		{1, 2, 1},
-		{1, 2, 2},
-		{1, 2, 3},
+func TestCombinationIndex(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
+		combs := make(map[string]struct{})
+		for i := 0; i < Binomial(n, k); i++ {
+			comb := IndexToCombination(nil, i, n, k)
+			idx := CombinationIndex(comb, n, k)
+			if idx != i {
+				t.Errorf("Cas %d: combination mismatch. Want %d, got %d", cas, i, idx)
+			}
+			combs[intSliceToKey(comb)] = struct{}{}
+		}
+		if len(combs) != Binomial(n, k) {
+			t.Errorf("Case %d: not all generated combinations were unique", cas)
+		}
+	}
+}
+
+func intSliceToKey(s []int) string {
+	var str string
+	for _, v := range s {
+		str += strconv.Itoa(v) + "_"
+	}
+	return str
+}
+
+// TestCombinationOrder tests that the different Combinations methods
+// agree on the iteration order.
+func TestCombinationOrder(t *testing.T) {
+	n := 7
+	k := 3
+	list := Combinations(n, k)
+	for i, v := range list {
+		idx := CombinationIndex(v, n, k)
+		if idx != i {
+			t.Errorf("Combinations and CombinationIndex mismatch")
+			break
 		}
-	got := Cartesian(lens)
-	if !intSosMatch(want, got) {
-		t.Errorf("cartesian data mismatch.\nwant:\n%v\ngot:\n%v", want, got)
 	}
 }
 
@@ -247,3 +262,38 @@ func TestIdxSubFor(t *testing.T) {
 		}
 	}
 }
+
+func TestCartesian(t *testing.T) {
+	// First, test with a known return.
+	lens := []int{2, 3, 4}
+	want := [][]int{
+		{0, 0, 0},
+		{0, 0, 1},
+		{0, 0, 2},
+		{0, 0, 3},
+		{0, 1, 0},
+		{0, 1, 1},
+		{0, 1, 2},
+		{0, 1, 3},
+		{0, 2, 0},
+		{0, 2, 1},
+		{0, 2, 2},
+		{0, 2, 3},
+		{1, 0, 0},
+		{1, 0, 1},
+		{1, 0, 2},
+		{1, 0, 3},
+		{1, 1, 0},
+		{1, 1, 1},
+		{1, 1, 2},
+		{1, 1, 3},
+		{1, 2, 0},
+		{1, 2, 1},
+		{1, 2, 2},
+		{1, 2, 3},
+	}
+	got := Cartesian(lens)
+	if !intSosMatch(want, got) {
+		t.Errorf("cartesian data mismatch.\nwant:\n%v\ngot:\n%v", want, got)
+	}
+}

stat/combin/combinations_example_test.go

@@ -0,0 +1,130 @@
+// Copyright ©2019 The Gonum Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package combin_test
+
+import (
+	"fmt"
+
+	"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.
+	fmt.Println("Generate list:")
+	n := 5
+	k := 3
+	list := combin.Combinations(n, k)
+	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.
+
+	// Output:
+	// Generate list:
+	// 0 [0 1 2]
+	// 1 [0 1 3]
+	// 2 [0 1 4]
+	// 3 [0 2 3]
+	// 4 [0 2 4]
+	// 5 [0 3 4]
+	// 6 [1 2 3]
+	// 7 [1 2 4]
+	// 8 [1 3 4]
+	// 9 [2 3 4]
+	//
+	// String combinations:
+	// ab
+	// ac
+	// ad
+	// ae
+	// bc
+	// bd
+	// be
+	// cd
+	// ce
+	// de
+
+}
+
+func ExampleCombinationGenerator() {
+	// combin provides several ways to work with the combinations of
+	// different objects. CombinationGenerator constructs an iterator
+	// for the combinations.
+	n := 5
+	k := 3
+	gen := combin.NewCombinationGenerator(n, k)
+	idx := 0
+	for gen.Next() {
+		fmt.Println(idx, gen.Combination(nil)) // can also store in-place.
+		idx++
+	}
+	// Output:
+	// 0 [0 1 2]
+	// 1 [0 1 3]
+	// 2 [0 1 4]
+	// 3 [0 2 3]
+	// 4 [0 2 4]
+	// 5 [0 3 4]
+	// 6 [1 2 3]
+	// 7 [1 2 4]
+	// 8 [1 3 4]
+	// 9 [2 3 4]
+}
+
+func ExampleIndexToCombination() {
+	// combin provides several ways to work with the combinations of
+	// different objects. IndexToCombination allows random access into
+	// the combination order. Combined with CombinationIndex this
+	// provides a correspondence between integers and combinations.
+	n := 5
+	k := 3
+	comb := make([]int, k)
+	for i := 0; i < combin.Binomial(n, k); i++ {
+		combin.IndexToCombination(comb, i, n, k) // can also use nil.
+		idx := combin.CombinationIndex(comb, n, k)
+		fmt.Println(i, comb, idx)
+	}
+
+	// Output:
+	// 0 [0 1 2] 0
+	// 1 [0 1 3] 1
+	// 2 [0 1 4] 2
+	// 3 [0 2 3] 3
+	// 4 [0 2 4] 4
+	// 5 [0 3 4] 5
+	// 6 [1 2 3] 6
+	// 7 [1 2 4] 7
+	// 8 [1 3 4] 8
+	// 9 [2 3 4] 9
+}