stats/combin: add CartesianGenerator (#1111)

* stats/combin: add CartesianGenerator

stat/combin/combin.go

@@ -284,23 +284,67 @@ func IndexToCombination(dst []int, idx, n, k int) []int {
 //  [ 1 2 0 ]
 // Cartesian panics if any of the provided lengths are less than 1.
 func Cartesian(lens []int) [][]int {
-	if len(lens) == 0 {
+	rows := Card(lens)
+	if rows == 0 {
 		panic("combin: empty lengths")
 	}
-	rows := 1
-	for _, v := range lens {
-		if v < 1 {
-			panic("combin: length less than zero")
-		}
-		rows *= v
-	}
 	out := make([][]int, rows)
 	for i := 0; i < rows; i++ {
 		out[i] = SubFor(nil, i, lens)
 	}
 	return out
 }
 
+// Card computes the cardinality of the multi-dimensional space whose dimensions have size specified by dims
+// All length values must be positive, otherwise this will panic.
+func Card(dims []int) int {
+	if len(dims) == 0 {
+		return 0
+	}
+	card := 1
+	for _, v := range dims {
+		if v < 1 {
+			panic("combin: length less than zero")
+		}
+		card *= v
+	}
+	return card
+}
+
+// NewCartesianGenerator returns a CartesianGenerator for iterating over cartesian products which are generated on the fly.
+// All values in lens must be positive, otherwise this will panic.
+func NewCartesianGenerator(lens []int) *CartesianGenerator {
+	return &CartesianGenerator{
+		lens: lens,
+		rows: Card(lens),
+		idx:  -1,
+	}
+}
+
+// CartesianGenerator iterates over a cartesian product set.
+type CartesianGenerator struct {
+	lens []int
+	rows int
+	idx  int
+}
+
+// Next moves to the next product of the cartesian set.
+// It returns false if the generator reached the end of the cartesian set end.
+func (g *CartesianGenerator) Next() bool {
+	if g.idx+1 < g.rows {
+		g.idx++
+		return true
+	}
+	g.idx = g.rows
+	return false
+}
+
+// Product generates one product of the cartesian set according to the current index which is increased by Next().
+// Next needs to be called at least one time before this method, otherwise it will panic.
+func (g *CartesianGenerator) Product() []int {
+	return SubFor(nil, g.idx, g.lens)
+}
+
 // IdxFor converts a multi-dimensional index into a linear index for a
 // multi-dimensional space. sub specifies the index for each dimension, and dims
 // specifies the size of each dimension. IdxFor is the inverse of SubFor.

stat/combin/combin_test.go

@@ -298,6 +298,38 @@ func TestCartesian(t *testing.T) {
 	}
 }
 
+func TestNumCartesianProducts(t *testing.T) {
+	want := 6
+	got := Card([]int{1, 2, 3})
+	if want != got {
+		t.Errorf("number of cartesian products mismatch.\nwant:\n%v\ngot:\n%v", want, got)
+	}
+}
+
+func TestCartesianGenerator(t *testing.T) {
+	want := [][]int{
+		{0, 0, 0},
+		{0, 0, 1},
+		{0, 0, 2},
+		{0, 1, 0},
+		{0, 1, 1},
+		{0, 1, 2},
+	}
+	gen := NewCartesianGenerator([]int{1, 2, 3})
+	iterations := 0
+	for gen.Next() {
+		got := gen.Product()
+		if !reflect.DeepEqual(got, want[iterations]) {
+			t.Errorf("Cartesian product does not match. want: %v got: %v", want[iterations], got)
+		}
+		iterations++
+	}
+
+	if iterations != len(want) {
+		t.Errorf("Number of products does not match. want: %v got: %v", len(want), iterations)
+	}
+}
+
 func TestPermutationIndex(t *testing.T) {
 	for cas, s := range []struct {
 		n, k int

stat/combin/combinations_example_test.go

@@ -10,6 +10,49 @@ import (
 	"gonum.org/v1/gonum/stat/combin"
 )
 
+func ExampleCartesian() {
+	fmt.Println("Generate cartesian products for given lengths:")
+	lens := []int{1, 2, 3}
+	list := combin.Cartesian(lens)
+	for i, v := range list {
+		fmt.Println(i, v)
+	}
+	// This is easy, but the number of combinations  can be very large,
+	// and generating all at once can use a lot of memory.
+	// For big data sets, consider using CartesianGenerator instead.
+
+	// Output:
+	// Generate cartesian products for given lengths:
+	// 0 [0 0 0]
+	// 1 [0 0 1]
+	// 2 [0 0 2]
+	// 3 [0 1 0]
+	// 4 [0 1 1]
+	// 5 [0 1 2]
+}
+
+func ExampleCartesianGenerator() {
+	fmt.Println("Generate products for given lengths:")
+	lens := []int{1, 2, 3}
+	gen := combin.NewCartesianGenerator(lens)
+
+	// Now loop over all products.
+	var i int
+	for gen.Next() {
+		fmt.Println(i, gen.Product())
+		i++
+	}
+
+	// Output:
+	// Generate products for given lengths:
+	// 0 [0 0 0]
+	// 1 [0 0 1]
+	// 2 [0 0 2]
+	// 3 [0 1 0]
+	// 4 [0 1 1]
+	// 5 [0 1 2]
+}
+
 func ExampleCombinations() {
 	// combin provides several ways to work with the combinations of
 	// different objects. Combinations generates them directly.