Generalize to any dimension, organizing in the process

2 files changed, 174 insertions, 167 deletions

diff --git a/capset.go b/capset.go
index 140c488..f28484a 100644
--- a/capset.go
+++ b/capset.go
@@ -4,65 +4,86 @@
 
 package main
 
-import "log"
-import "os"
-import "fmt"
-import "bufio"
-import "errors"
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"log"
+	"math"
+	"os"
+)
 
-var deck [81]string
-var card2IntMap map[string]int
-var useMap = true
+var maximForDim = []int{1, 2, 4, 9, 20, 45, 112, 236}
 
-func init() {
-	for i := range deck {
-		deck[i] = computeInt2Card(80 - i)
+type Card string
+
+type Deck struct {
+	// dim is the number of dimensions for the set.
+	dim int
+	// cards is the full deck of cards, sorted.
+	cards []Card
+	// card2index maps from a card back to its position in cards slice.
+	card2index map[Card]int
+}
+
+func NewDeck(dimensions int) *Deck {
+	size := int(math.Pow(3, float64(dimensions)))
+	d := &Deck{
+		dim:        dimensions,
+		cards:      make([]Card, size),
+		card2index: make(map[Card]int, size),
 	}
-	card2IntMap = make(map[string]int)
-	for i, v := range deck {
-		card2IntMap[v] = i
+	for i := range d.cards {
+		card := d.computeInt2Card(size - i - 1)
+		d.cards[i] = card
+		d.card2index[card] = i
 	}
+	return d
 }
 
-func computeInt2Card(num int) string {
+func (d *Deck) computeInt2Card(num int) Card {
 	if num < 0 {
-		log.Fatal("Number must be nonnegative.")
+		panic("Number must be positive")
 	}
-	if num > 81 {
-		log.Fatal("Number must be less than 81.")
-	}
-	var strlist [4]byte
-	for i := 0; i < 4; i++ {
-		strlist[3-i] = '0' + byte(num%3)
+	strlist := make([]byte, d.dim)
+	for i := 0; i < d.dim; i++ {
+		// Always have three options for each dimension
+		strlist[d.dim-i-1] = '0' + byte(num%3)
 		num /= 3
 	}
-	return string(strlist[:])
+	return Card(strlist)
 }
 
-func card2int(item string) int {
-	if useMap {
-		i, ok := card2IntMap[item]
-		if !ok {
-			return -1
-		}
-		return i
-	} else {
-		for i, v := range deck {
-			if v == item {
-				return i
-			}
-		}
+func (d *Deck) card2int(card Card) int {
+	i, ok := d.card2index[card]
+	if !ok {
 		return -1
 	}
+	return i
 }
 
-type capset struct {
-	tableau []string
+type Capset struct {
+	d       *Deck
+	tableau []Card
 	maxim   int
 	hasSet  bool
 }
 
-func (cs *capset) Scan(state fmt.ScanState, verb rune) error {
+// NewCapset makes an initialized Capset for the provided deck. If maxim is -1,
+// then the maximum known capset size for the deck's dimension will be used, or
+// 0 if it is unknown.
+func NewCapset(d *Deck, maxim int) *Capset {
+	if maxim == -1 {
+		maxim = maximForDim[d.dim]
+	}
+	return &Capset{
+		d:       d,
+		tableau: make([]Card, 0, len(d.cards)),
+		maxim:   maxim,
+	}
+}
+
+func (cs *Capset) Scan(state fmt.ScanState, verb rune) error {
 	if verb != 'v' {
 		return errors.New("unexpected verb")
 	}
@@ -87,165 +108,137 @@ func (cs *capset) Scan(state fmt.ScanState, verb rune) error {
 			token = token[:len(token)-1]
 		}
 		if len(token) > 0 {
-			card := string(token)
-			index := card2int(card)
+			card := Card(token)
+			index := cs.d.card2int(card)
 			if index == -1 {
 				return errors.New(fmt.Sprintf("Invalid card: %v", card))
 			}
 			// Always use same strings in hope underlying strcmp short circuits
-			card = deck[index]
+			card = cs.d.cards[index]
 			cs.tableau = append(cs.tableau, card)
 		}
 	}
-	cs.updateTableauHasSetCached()
+	cs.hasSet = hasSet(cs.tableau)
 	// Loop until it isn't a set, to make sure we didn't load a tableau with lots
 	// of sets.
 	for cs.hasSet {
-		cs.incrementTableau(true)
+		cs.incrementTableau()
+		cs.hasSet = hasSet(cs.tableau)
 	}
 	return nil
 }
 
-func (cs *capset) popTableau() string {
+func (cs *Capset) popTableau() Card {
 	v := cs.tableau[len(cs.tableau)-1]
 	cs.tableau = cs.tableau[:len(cs.tableau)-1]
 	return v
 }
 
-func (cs *capset) incrementTableau(noAssume bool) bool {
+func (cs *Capset) incrementTableau() bool {
 	pop := cs.hasSet
-	index := card2int(cs.tableau[len(cs.tableau)-1]) + 1
-	if index != len(deck) {
+	var index int
+	if len(cs.tableau) > 0 {
+		index = cs.d.card2int(cs.tableau[len(cs.tableau)-1]) + 1
+	} else {
+		index = 0
+	}
+	if index != len(cs.d.cards) {
 		if pop {
 			cs.popTableau()
 		}
-		cs.tableau = append(cs.tableau, deck[index])
-		cs.updateTableauHasSet(noAssume)
+		cs.tableau = append(cs.tableau, cs.d.cards[index])
 		return true
 	}
 	cs.popTableau()
-	index = card2int(cs.popTableau()) + 1
+	if len(cs.tableau) == 0 {
+		// No combinations left
+		return false
+	}
+	index = cs.d.card2int(cs.popTableau()) + 1
 	if len(cs.tableau) == 1 {
 		fmt.Println("incrementing second card to index", index)
-		if index+18 > len(deck) {
+		if index+18 > len(cs.d.cards) {
 			return false
 		}
 	}
-	cs.tableau = append(cs.tableau, deck[index])
-	cs.updateTableauHasSet(noAssume)
+	cs.tableau = append(cs.tableau, cs.d.cards[index])
 	return true
 }
 
-func isSet(cards [3]string) bool {
-	for i := 0; i < 4; i++ {
-		if cards[0][i] == cards[1][i] && cards[0][i] == cards[2][i] {
-			continue
-		}
-		if cards[0][i] != cards[1][i] &&
-			cards[0][i] != cards[2][i] &&
-			cards[1][i] != cards[2][i] {
-			continue
-		}
+// hasSet returns whether the provide tableau contains at least one set.
+func hasSet(tableau []Card) bool {
+	if len(tableau) <= 2 {
 		return false
 	}
-	return true
-}
-
-func (cs *capset) updateTableauHasSet(noAssume bool) {
-	if noAssume {
-		cs.updateTableauHasSetCached()
-	} else {
-		cs.updateTableauHasSetOnlyNew()
-	}
-}
-
-// TODO(ejona): remove
-func (cs *capset) updateTableauHasSetSimple() {
-	if len(cs.tableau) <= 2 {
-		cs.hasSet = false
-		return
-	}
-	var comb [3]string
-	for i := 0; i < len(cs.tableau); i++ {
-		comb[0] = cs.tableau[i]
-		for j := i + 1; j < len(cs.tableau); j++ {
-			comb[1] = cs.tableau[j]
-			for k := j + 1; k < len(cs.tableau); k++ {
-				comb[2] = cs.tableau[k]
-				if isSet(comb) {
-					cs.hasSet = true
-					return
-				}
-			}
-		}
-	}
-	cs.hasSet = false
-}
-
-func (cs *capset) updateTableauHasSetCached() {
-	var same [4]bool
-	for i := 0; i < len(cs.tableau); i++ {
-		for j := i + 1; j < len(cs.tableau); j++ {
-			for x := 0; x < 4; x++ {
-				same[x] = cs.tableau[i][x] == cs.tableau[j][x]
+	dim := len(tableau[0])
+	same := make([]bool, dim)
+	for i := 0; i < len(tableau); i++ {
+		for j := i + 1; j < len(tableau); j++ {
+			for x := 0; x < dim; x++ {
+				same[x] = tableau[i][x] == tableau[j][x]
 			}
 		kloop:
-			for k := j + 1; k < len(cs.tableau); k++ {
-				for x := 0; x < 4; x++ {
+			for k := j + 1; k < len(tableau); k++ {
+				for x := 0; x < dim; x++ {
 					if same[x] {
-						if cs.tableau[i][x] != cs.tableau[k][x] {
+						if tableau[i][x] != tableau[k][x] {
 							continue kloop
 						}
 					} else {
-						if cs.tableau[i][x] == cs.tableau[k][x] ||
-							cs.tableau[j][x] == cs.tableau[k][x] {
+						if tableau[i][x] == tableau[k][x] ||
+							tableau[j][x] == tableau[k][x] {
 							continue kloop
 						}
 					}
 				}
-				cs.hasSet = true
-				return
+				return true
 			}
 		}
 	}
-	cs.hasSet = false
+	return false
 }
 
-func (cs *capset) updateTableauHasSetOnlyNew() {
-	if len(cs.tableau) <= 2 {
-		cs.hasSet = false
-		return
+// hasSetWithLast returns whether the provided tableau contains at least one set
+// with the last element of tableau as a member.
+func hasSetWithLast(tableau []Card) bool {
+	if len(tableau) <= 2 {
+		return false
 	}
 	// We try sets such that only the rightmost item could possibly be in a set.
-	i := len(cs.tableau) - 1
-	var same [4]bool
+	i := len(tableau) - 1
+	dim := len(tableau[0])
+	// Using make instead of "var same [4]bool" causes a ~6% slowdown
+	same := make([]bool, dim)
 	for j := 0; j < i; j++ {
-		for x := 0; x < 4; x++ {
-			same[x] = cs.tableau[i][x] == cs.tableau[j][x]
+		for x := 0; x < dim; x++ {
+			same[x] = tableau[i][x] == tableau[j][x]
 		}
 	kloop:
 		for k := j + 1; k < i; k++ {
-			for x := 0; x < 4; x++ {
+			for x := 0; x < dim; x++ {
 				if same[x] {
-					if cs.tableau[i][x] != cs.tableau[k][x] {
+					if tableau[i][x] != tableau[k][x] {
 						continue kloop
 					}
 				} else {
-					if cs.tableau[i][x] == cs.tableau[k][x] ||
-						cs.tableau[j][x] == cs.tableau[k][x] {
+					if tableau[i][x] == tableau[k][x] ||
+						tableau[j][x] == tableau[k][x] {
 						continue kloop
 					}
 				}
 			}
-			cs.hasSet = true
-			return
+			return true
 		}
 	}
-	cs.hasSet = false
+	return false
 }
 
-func (cs *capset) findNextCapset() bool {
-	for cs.incrementTableau(false) {
+// FindNextCapset iterates over combinations of the tableau until it finds a
+// capset that is at least the maxim in size. It returns false if it was unable
+// to find such a capset.
+func (cs *Capset) FindNextCapset() bool {
+	for cs.incrementTableau() {
+		cs.hasSet = hasSetWithLast(cs.tableau)
 		if cs.hasSet {
 			continue
 		}
@@ -258,7 +251,12 @@ func (cs *capset) findNextCapset() bool {
 	return false
 }
 
-func loadFromSave() *capset {
+// GetTableau returns the current tableau. Callers must not modify the slice.
+func (cs *Capset) GetTableau() []Card {
+	return cs.tableau
+}
+
+func loadFromSave(d *Deck) *Capset {
 	savefile, err := os.Open("./capset.out")
 	if err != nil {
 		return nil
@@ -267,20 +265,20 @@ func loadFromSave() *capset {
 	scanner := bufio.NewScanner(savefile)
 	for scanner.Scan() {
 	}
-	var cs capset
-	_, err = fmt.Sscanln(scanner.Text(), &cs)
+	cs := NewCapset(d, -1)
+	_, err = fmt.Sscanln(scanner.Text(), cs)
 	if err != nil {
 		return nil
 	}
-	return &cs
+	return cs
 }
 
 func main() {
-	cs := loadFromSave()
+	d := NewDeck(4)
+	cs := loadFromSave(d)
 	if cs == nil {
-		cs = &capset{tableau: []string{"2222"}}
+		cs = NewCapset(d, -1)
 	}
-	cs.maxim = 20
 	fmt.Println("Starting processing with", cs.tableau)
 
 	outfile, err := os.OpenFile("./capset.out",
@@ -290,9 +288,9 @@ func main() {
 	}
 	defer outfile.Close()
 
-	for cs.findNextCapset() {
-		fmt.Println(cs.tableau)
-		_, err = fmt.Fprintln(outfile, cs.tableau)
+	for cs.FindNextCapset() {
+		fmt.Println(cs.GetTableau())
+		_, err = fmt.Fprintln(outfile, cs.GetTableau())
 		if err != nil {
 			log.Println(err)
 		}
diff --git a/capset_test.go b/capset_test.go
index 5382aa8..cd2715e 100644
--- a/capset_test.go
+++ b/capset_test.go
@@ -1,25 +1,29 @@
 package main
 
-import "testing"
-import "reflect"
-import "fmt"
+import (
+	"fmt"
+	"reflect"
+	"testing"
+)
+
+var deck = NewDeck(4)
 
 var capsettests = []struct {
 	maxim         int
-	start, finish []string
+	start, finish []Card
 }{
-	{20, []string{"2222", "2221", "2212", "2211", "2122", "2121", "2112", "2111",
+	{20, []Card{"2222", "2221", "2212", "2211", "2122", "2121", "2112", "2111",
 		"1222", "1221", "1210"},
-		[]string{"2222", "2221", "2212", "2211", "2122", "2121", "2112", "2111",
+		[]Card{"2222", "2221", "2212", "2211", "2122", "2121", "2112", "2111",
 			"1222", "1221", "1210", "1200", "1120", "1020", "0210", "0120", "0112",
 			"0111", "0100", "0010"}},
 }
 
 func TestFindNextCapset(t *testing.T) {
 	for _, tt := range capsettests {
-		cs := capset{tableau: tt.start, maxim: tt.maxim}
-		cs.findNextCapset()
-		expected := capset{tableau: tt.finish, maxim: len(tt.finish)}
+		cs := Capset{tableau: tt.start, maxim: tt.maxim, d: deck}
+		cs.FindNextCapset()
+		expected := Capset{tableau: tt.finish, maxim: len(tt.finish), d: deck}
 		if !reflect.DeepEqual(cs, expected) {
 			t.Fatalf("Got \n%v but expected \n%v", cs, expected)
 		}
@@ -28,26 +32,26 @@ func TestFindNextCapset(t *testing.T) {
 
 var scantests = []struct {
 	in  string
-	out []string
+	out []Card
 	err bool
 }{
 	{"", nil, true},
-	{"[]", []string{}, false},
-	{"[0000]", []string{"0000"}, false},
-	{"[0000 2222]", []string{"0000", "2222"}, false},
+	{"[]", []Card{}, false},
+	{"[0000]", []Card{"0000"}, false},
+	{"[0000 2222]", []Card{"0000", "2222"}, false},
 	{"[2]", nil, true},
 	{"not it", nil, true},
 	{"[0000", nil, true},
-	{"[  0000   2222  ]", []string{"0000", "2222"}, false},
+	{"[  0000   2222  ]", []Card{"0000", "2222"}, false},
 }
 
 func TestScan(t *testing.T) {
 	for _, tt := range scantests {
-		cs := capset{tableau: []string{}}
+		cs := NewCapset(deck, -1)
 		if testing.Verbose() {
 			fmt.Println("Trying input", tt.in)
 		}
-		_, err := fmt.Sscanln(tt.in, &cs)
+		_, err := fmt.Sscanln(tt.in, cs)
 		if tt.err {
 			if err == nil {
 				t.Errorf("Sscanln(%v) => %v, want an error", tt.in, cs.tableau)
@@ -64,28 +68,33 @@ func TestScan(t *testing.T) {
 	}
 }
 
-func testNoAllocs(t *testing.T) {
-	cs := capset{tableau: []string{"2222"}}
-	result := testing.AllocsPerRun(10, func() { cs.incrementTableau(false) })
+func TestNoAllocs(t *testing.T) {
+	cs := NewCapset(deck, -1)
+	result := testing.AllocsPerRun(10, func() { cs.incrementTableau() })
 	if result > 0 {
 		t.Errorf("Too many allocs: %v", result)
 	}
 }
 
-func BenchmarkFindNextCapset(b *testing.B) {
-	cs := capset{tableau: []string{"2222"}}
-	useMap = true
+func BenchmarkIncrement(b *testing.B) {
+	cs := NewCapset(deck, -1)
 	for i := 0; i < b.N; i++ {
-		cs.incrementTableau(false)
+		cs.incrementTableau()
 	}
 }
 
-func BenchmarkFindNextCapsetNoMap(b *testing.B) {
-	cs := capset{tableau: []string{"2222"}}
-	useMap = false
+func BenchmarkIncrementHasSet(b *testing.B) {
+	cs := NewCapset(deck, -1)
 	for i := 0; i < b.N; i++ {
-		cs.incrementTableau(false)
+		cs.incrementTableau()
+		cs.hasSet = hasSet(cs.tableau)
 	}
-	useMap = true
 }
 
+func BenchmarkIncrementHasSetWithLast(b *testing.B) {
+	cs := NewCapset(deck, -1)
+	for i := 0; i < b.N; i++ {
+		cs.incrementTableau()
+		cs.hasSet = hasSetWithLast(cs.tableau)
+	}
+}