- add reversible iterators to red-black tree

trees/redblacktree/redblacktree.go

@@ -43,7 +43,7 @@ import (
 
 func assertInterfaceImplementation() {
 	var _ trees.Tree = (*Tree)(nil)
-	var _ containers.IteratorWithKey = (*Iterator)(nil)
+	var _ containers.ReverseIteratorWithKey = (*Iterator)(nil)
 }
 
 type color bool
@@ -281,37 +281,65 @@ func (tree *Tree) Clear() {
 // Iterator holding the iterator's state
 type Iterator struct {
 	tree *Tree
-	left *Node
+	node *Node
 }
 
 // Iterator returns a stateful iterator whose elements are key/value pairs.
 func (tree *Tree) Iterator() Iterator {
-	return Iterator{tree: tree, left: nil}
+	return Iterator{tree: tree, node: nil}
 }
 
 // Next moves the iterator to the next element and returns true if there was a next element in the container.
 // If Next() returns true, then next element's key and value can be retrieved by Key() and Value().
 // Modifies the state of the iterator.
 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
+	if iterator.node == nil {
+		iterator.node = iterator.tree.Left()
+		return iterator.node != nil
+	}
+	if iterator.node.Right != nil {
+		iterator.node = iterator.node.Right
+		for iterator.node.Left != nil {
+			iterator.node = iterator.node.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 {
+	if iterator.node.Parent != nil {
+		node := iterator.node
+		for iterator.node.Parent != nil {
+			iterator.node = iterator.node.Parent
+			if iterator.tree.Comparator(node.Key, iterator.node.Key) <= 0 {
 				return true
 			}
 		}
+		iterator.node = node // fix: if parent didn't satisfy the comparator criteria
+	}
+	return false
+}
+
+// Prev moves the iterator to the previous element and returns true if there was a previous element in the container.
+// If Prev() returns true, then previous element's index and value can be retrieved by Key() and Value().
+// Modifies the state of the iterator.
+func (iterator *Iterator) Prev() bool {
+	if iterator.node == nil {
+		return false
+	}
+	if iterator.node.Left != nil {
+		iterator.node = iterator.node.Left
+		for iterator.node.Right != nil {
+			iterator.node = iterator.node.Right
+		}
+		return true
+	}
+	if iterator.node.Parent != nil {
+		node := iterator.node
+		for iterator.node.Parent != nil {
+			iterator.node = iterator.node.Parent
+			if iterator.tree.Comparator(node.Key, iterator.node.Key) >= 0 {
+				return true
+			}
+		}
+		iterator.node = node // fix: if parent didn't satisfy the comparator criteria
 	}
 	return false
 }
@@ -319,13 +347,13 @@ func (iterator *Iterator) Next() bool {
 // Value returns the current element's value.
 // Does not modify the state of the iterator.
 func (iterator *Iterator) Value() interface{} {
-	return iterator.left.Value
+	return iterator.node.Value
 }
 
 // Key returns the current element's key.
 // Does not modify the state of the iterator.
 func (iterator *Iterator) Key() interface{} {
-	return iterator.left.Key
+	return iterator.node.Key
 }
 
 // String returns a string representation of container

trees/redblacktree/redblacktree_test.go

@@ -206,7 +206,23 @@ func TestRedBlackTreeCeilingAndFloor(t *testing.T) {
 	}
 }
 
-func TestRedBlackTreeIterator1(t *testing.T) {
+func TestRedBlackTreeIteratorNextOnEmpty(t *testing.T) {
+	tree := NewWithIntComparator()
+	it := tree.Iterator()
+	for it.Next() {
+		t.Errorf("Shouldn't iterate on empty tree")
+	}
+}
+
+func TestRedBlackTreeIteratorPrevOnEmpty(t *testing.T) {
+	tree := NewWithIntComparator()
+	it := tree.Iterator()
+	for it.Prev() {
+		t.Errorf("Shouldn't iterate on empty tree")
+	}
+}
+
+func TestRedBlackTreeIterator1Next(t *testing.T) {
 	tree := NewWithIntComparator()
 	tree.Put(5, "e")
 	tree.Put(6, "f")
@@ -216,29 +232,76 @@ func TestRedBlackTreeIterator1(t *testing.T) {
 	tree.Put(1, "x")
 	tree.Put(2, "b")
 	tree.Put(1, "a") //overwrite
-
+	// │   ┌── 7
+	// └── 6
+	//     │   ┌── 5
+	//     └── 4
+	//         │   ┌── 3
+	//         └── 2
+	//             └── 1
 	it := tree.Iterator()
 	count := 0
 	for it.Next() {
 		count++
-		index := it.Key()
-		switch index {
+		key := it.Key()
+		switch key {
 		case count:
-			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+			if actualValue, expectedValue := key, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			if actualValue, expectedValue := key, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		}
+	}
+	if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue {
+		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
+	}
+}
+
+func TestRedBlackTreeIterator1Prev(t *testing.T) {
+	tree := NewWithIntComparator()
+	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
+	// │   ┌── 7
+	// └── 6
+	//     │   ┌── 5
+	//     └── 4
+	//         │   ┌── 3
+	//         └── 2
+	//             └── 1
+	it := tree.Iterator()
+	for it.Next() {
+	}
+	countDown := tree.size
+	for it.Prev() {
+		countDown--
+		key := it.Key()
+		switch key {
+		case countDown:
+			if actualValue, expectedValue := key, countDown; actualValue != expectedValue {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		default:
-			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+			if actualValue, expectedValue := key, countDown; actualValue != expectedValue {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		}
 	}
-	if actualValue, expectedValue := count, 7; actualValue != expectedValue {
+	// one less that in Next(), thus "1"
+	if actualValue, expectedValue := countDown, 1; actualValue != expectedValue {
 		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
 	}
 }
 
-func TestRedBlackTreeIterator2(t *testing.T) {
+func TestRedBlackTreeIterator2Next(t *testing.T) {
 	tree := NewWithIntComparator()
 	tree.Put(3, "c")
 	tree.Put(1, "a")
@@ -247,50 +310,99 @@ func TestRedBlackTreeIterator2(t *testing.T) {
 	count := 0
 	for it.Next() {
 		count++
-		index := it.Key()
-		switch index {
+		key := it.Key()
+		switch key {
 		case count:
-			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+			if actualValue, expectedValue := key, count; actualValue != expectedValue {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		default:
-			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+			if actualValue, expectedValue := key, count; actualValue != expectedValue {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		}
 	}
-	if actualValue, expectedValue := count, 3; actualValue != expectedValue {
+	if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue {
 		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
 	}
 }
 
-func TestRedBlackTreeIterator3(t *testing.T) {
+func TestRedBlackTreeIterator2Prev(t *testing.T) {
 	tree := NewWithIntComparator()
-
+	tree.Put(3, "c")
+	tree.Put(1, "a")
+	tree.Put(2, "b")
 	it := tree.Iterator()
 	for it.Next() {
-		t.Errorf("Shouldn't iterate on empty stack")
 	}
+	countDown := tree.size
+	for it.Prev() {
+		countDown--
+		key := it.Key()
+		switch key {
+		case countDown:
+			if actualValue, expectedValue := key, countDown; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			if actualValue, expectedValue := key, countDown; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		}
+	}
+	// one less that in Next(), thus "1"
+	if actualValue, expectedValue := countDown, 1; actualValue != expectedValue {
+		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
+	}
+}
 
+func TestRedBlackTreeIterator3Next(t *testing.T) {
+	tree := NewWithIntComparator()
 	tree.Put(1, "a")
-
-	it = tree.Iterator()
+	it := tree.Iterator()
 	count := 0
 	for it.Next() {
 		count++
-		index := it.Key()
-		switch index {
+		key := it.Key()
+		switch key {
 		case count:
-			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+			if actualValue, expectedValue := key, count; actualValue != expectedValue {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		default:
-			if actualValue, expectedValue := index, count; actualValue != expectedValue {
+			if actualValue, expectedValue := key, count; actualValue != expectedValue {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		}
 	}
-	if actualValue, expectedValue := count, 1; actualValue != expectedValue {
+	if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue {
+		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
+	}
+}
+
+func TestRedBlackTreeIterator3Prev(t *testing.T) {
+	tree := NewWithIntComparator()
+	tree.Put(1, "a")
+	it := tree.Iterator()
+	for it.Next() {
+	}
+	countDown := tree.size
+	for it.Prev() {
+		countDown--
+		key := it.Key()
+		switch key {
+		case countDown:
+			if actualValue, expectedValue := key, countDown; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			if actualValue, expectedValue := key, countDown; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		}
+	}
+	// one less that in Next(), thus "1"
+	if actualValue, expectedValue := countDown, 1; actualValue != expectedValue {
 		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
 	}
 }
@@ -308,7 +420,20 @@ func TestRedBlackTreeIterator4(t *testing.T) {
 	tree.Put(22, 8)
 	tree.Put(27, 10)
 
+	// │           ┌── 27
+	// │       ┌── 25
+	// │       │   └── 22
+	// │   ┌── 17
+	// │   │   └── 15
+	// └── 13
+	//     │   ┌── 11
+	//     └── 8
+	//         │   ┌── 6
+	//         └── 1
+
 	it := tree.Iterator()
+
+	// Iterator (next)
 	count := 0
 	for it.Next() {
 		count++
@@ -324,7 +449,27 @@ func TestRedBlackTreeIterator4(t *testing.T) {
 			}
 		}
 	}
-	if actualValue, expectedValue := count, 10; actualValue != expectedValue {
+	if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue {
+		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
+	}
+
+	// Iterator (prev)
+	for it.Prev() {
+		count--
+		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)
+			}
+		}
+	}
+	// one less that in Next(), thus "1"
+	if actualValue, expectedValue := count, 1; actualValue != expectedValue {
 		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
 	}
 }