1 files changed, 36 insertions, 12 deletions
diff --git a/set/set_test.go b/set/set_test.go
index dac258a..9c5f102 100644
--- a/set/set_test.go
+++ b/set/set_test.go
@@ -5,11 +5,12 @@ import (
 	"testing"
 
 	"github.com/snonux/algorithms/ds"
+	"github.com/snonux/algorithms/sort"
 )
 
 const factor int = 10
 const minLength int = 1
-const maxLength int = 10000
+const maxLength int = 10
 
 // Store results here to avoid compiler optimizations
 var benchResult int
@@ -22,29 +23,45 @@ func TestElementary(t *testing.T) {
 }
 
 func TestBST(t *testing.T) {
-	s := NewBST()
 	for i := minLength; i <= maxLength; i *= factor {
-		test(s, i, t)
+		test(NewBST(), i, t)
+	}
+}
+
+func TestBalancedBST(t *testing.T) {
+	for i := minLength; i <= maxLength; i *= factor {
+		test(NewBalancedBST(), i, t)
 	}
 }
 
 func test(s Set, l int, t *testing.T) {
 	cb := func(t *testing.T) {
-		list := ds.NewRandomArrayList(l, -1)
-		for i, a := range list {
-			t.Log("Inserting", a, i, s)
-			s.Set(a, i)
+		vals := ds.NewRandomArrayList(l, -1)
+		keys := ds.NewRandomArrayList(l, -1)
+		mapping := make(map[int]int, l)
+
+		for i, key := range keys {
+			val := vals[i]
+			mapping[key] = val
+
+			t.Log("Inserting", key, val)
+			s.Set(key, val)
 		}
-		for i, a := range list {
-			val, err := s.Get(a)
+
+		for _, key := range sort.Shuffle(keys) {
+			val, err := s.Get(key)
 			if err != nil {
-				t.Errorf("Element %v: %v", val, err)
+				t.Errorf("Element %v->%v: %v\n", key, val, err)
 			}
-			if val != i {
-				t.Errorf("Element is %v but expected %v", val, i)
+
+			val2 := mapping[key]
+			if val2 != val {
+				t.Errorf("Element is %v->%v but expected %v\n", key, val, val2)
 			}
+			t.Log("Got", key, val, s)
 		}
 	}
+
 	t.Run(fmt.Sprintf("%d", l), cb)
 }
 
@@ -62,6 +79,13 @@ func BenchmarkBST(t *testing.B) {
 	}
 }
 
+func BenchmarkBalancedBST(t *testing.B) {
+	s := NewBalancedBST()
+	for i := minLength; i <= maxLength; i *= factor {
+		benchmark(s, i, t)
+	}
+}
+
 func benchmark(s Set, l int, b *testing.B) {
 	list := ds.NewRandomArrayList(l, -1)