diff --git a/trees/redblacktree/redblacktree.go b/trees/redblacktree/redblacktree.go index adcd4ec..cf68bab 100644 --- a/trees/redblacktree/redblacktree.go +++ b/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 diff --git a/trees/redblacktree/redblacktree_test.go b/trees/redblacktree/redblacktree_test.go index 96fa4e1..377ee9e 100644 --- a/trees/redblacktree/redblacktree_test.go +++ b/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) } }