- add iterator to redblack tree

trees/redblacktree/redblacktree.go

@@ -33,6 +33,7 @@ package redblacktree
 
 import (
 	"fmt"
+	"github.com/emirpasic/gods/containers"
 	"github.com/emirpasic/gods/stacks/linkedliststack"
 	"github.com/emirpasic/gods/trees"
 	"github.com/emirpasic/gods/utils"
@@ -40,6 +41,7 @@ import (
 
 func assertInterfaceImplementation() {
 	var _ trees.Tree = (*Tree)(nil)
+	var _ containers.Iterator = (*Iterator)(nil)
 }
 
 type color bool
@@ -272,6 +274,47 @@ func (tree *Tree) Clear() {
 	tree.size = 0
 }
 
+type Iterator struct {
+	tree *Tree
+	left *Node
+}
+
+func (tree *Tree) Iterator() Iterator {
+	return Iterator{tree: tree, left: nil}
+}
+
+func (iterator *Iterator) Next() bool {
+	if iterator.left == nil {
+		iterator.left = iterator.tree.Left()
+		return iterator.left != nil
+	}
+	if iterator.left.Right != nil {
+		iterator.left = iterator.left.Right
+		for iterator.left.Left != nil {
+			iterator.left = iterator.left.Left
+		}
+		return true
+	}
+	if iterator.left.Parent != nil {
+		key := iterator.left.Key
+		for iterator.left.Parent != nil {
+			iterator.left = iterator.left.Parent
+			if iterator.tree.Comparator(key, iterator.left.Key) <= 0 {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+func (iterator *Iterator) Value() interface{} {
+	return iterator.left.Value
+}
+
+func (iterator *Iterator) Index() interface{} {
+	return iterator.left.Key
+}
+
 func (tree *Tree) String() string {
 	str := "RedBlackTree\n"
 	if !tree.Empty() {

trees/redblacktree/redblacktree_test.go

@@ -216,6 +216,128 @@ func TestRedBlackTree(t *testing.T) {
 	}
 }
 
+func TestRedBlackTreeIterator(t *testing.T) {
+	tree := NewWithIntComparator()
+
+	// insertions
+	tree.Put(5, "e")
+	tree.Put(6, "f")
+	tree.Put(7, "g")
+	tree.Put(3, "c")
+	tree.Put(4, "d")
+	tree.Put(1, "x")
+	tree.Put(2, "b")
+	tree.Put(1, "a") //overwrite
+
+	// Iterator
+	it := tree.Iterator()
+	count := 0
+	for it.Next() {
+		count += 1
+		index := it.Index()
+		switch index {
+		case count:
+			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		}
+	}
+	if actualValue, expectedValue := count, 7; actualValue != expectedValue {
+		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
+	}
+
+	// Iterator
+	tree.Clear()
+	tree.Put(3, "c")
+	tree.Put(1, "a")
+	tree.Put(2, "b")
+	it = tree.Iterator()
+	count = 0
+	for it.Next() {
+		count += 1
+		index := it.Index()
+		switch index {
+		case count:
+			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		}
+	}
+	if actualValue, expectedValue := count, 3; actualValue != expectedValue {
+		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
+	}
+
+	// Iterator
+	tree.Clear()
+	tree.Put(1, "a")
+	it = tree.Iterator()
+	count = 0
+	for it.Next() {
+		count += 1
+		index := it.Index()
+		switch index {
+		case count:
+			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		}
+	}
+	if actualValue, expectedValue := count, 1; actualValue != expectedValue {
+		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
+	}
+
+	// Iterator on empty
+	tree.Clear()
+	it = tree.Iterator()
+	for it.Next() {
+		t.Errorf("Shouldn't iterate on empty stack")
+	}
+
+	// Iterator (from image)
+	tree.Clear()
+	tree.Put(13, 5)
+	tree.Put(8, 3)
+	tree.Put(17, 7)
+	tree.Put(1, 1)
+	tree.Put(11, 4)
+	tree.Put(15, 6)
+	tree.Put(25, 9)
+	tree.Put(6, 2)
+	tree.Put(22, 8)
+	tree.Put(27, 10)
+	it = tree.Iterator()
+	count = 0
+	for it.Next() {
+		count += 1
+		value := it.Value()
+		switch value {
+		case count:
+			if actualValue, expectedValue := value, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			if actualValue, expectedValue := value, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		}
+	}
+	if actualValue, expectedValue := count, 10; actualValue != expectedValue {
+		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
+	}
+}
+
 func BenchmarkRedBlackTree(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		tree := NewWithIntComparator()