- add reversible iterators to tree set and tree map

emirpasic · Jun 25, 2016 · eb4bb22 · eb4bb22
1 parent 178bc76
commit eb4bb22
Show file tree

Hide file tree

Showing 5 changed files with 158 additions and 19 deletions.
diff --git a/maps/treemap/treemap.go b/maps/treemap/treemap.go
@@ -43,7 +43,7 @@ import (
 func assertInterfaceImplementation() {
 	var _ maps.Map = (*Map)(nil)
 	var _ containers.EnumerableWithKey = (*Map)(nil)
-	var _ containers.IteratorWithKey = (*Iterator)(nil)
+	var _ containers.ReverseIteratorWithKey = (*Iterator)(nil)
 }
 
 // Map holds the elements in a red-black tree
@@ -145,6 +145,13 @@ func (iterator *Iterator) Next() bool {
 	return iterator.iterator.Next()
 }
 
+// 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 {
+	return iterator.iterator.Prev()
+}
+
 // Value returns the current element's value.
 // Does not modify the state of the iterator.
 func (iterator *Iterator) Value() interface{} {

diff --git a/maps/treemap/treemap_test.go b/maps/treemap/treemap_test.go
@@ -311,7 +311,25 @@ func TestMapChaining(t *testing.T) {
 	}
 }
 
-func TestMapIterator(t *testing.T) {
+func TestMapIteratorNextOnEmpty(t *testing.T) {
+	m := NewWithStringComparator()
+	it := m.Iterator()
+	it = m.Iterator()
+	for it.Next() {
+		t.Errorf("Shouldn't iterate on empty map")
+	}
+}
+
+func TestMapIteratorPrevOnEmpty(t *testing.T) {
+	m := NewWithStringComparator()
+	it := m.Iterator()
+	it = m.Iterator()
+	for it.Prev() {
+		t.Errorf("Shouldn't iterate on empty map")
+	}
+}
+
+func TestMapIteratorNext(t *testing.T) {
 	m := NewWithStringComparator()
 	m.Put("c", 3)
 	m.Put("a", 1)
@@ -346,11 +364,45 @@ func TestMapIterator(t *testing.T) {
 	if actualValue, expectedValue := count, 3; actualValue != expectedValue {
 		t.Errorf("Got %v expected %v", actualValue, expectedValue)
 	}
+}
 
-	m.Clear()
-	it = m.Iterator()
+func TestMapIteratorPrev(t *testing.T) {
+	m := NewWithStringComparator()
+	m.Put("c", 3)
+	m.Put("a", 1)
+	m.Put("b", 2)
+
+	it := m.Iterator()
 	for it.Next() {
-		t.Errorf("Shouldn't iterate on empty map")
+	}
+	countDown := m.Size()
+	for it.Prev() {
+		countDown--
+		key := it.Key()
+		value := it.Value()
+		switch key {
+		case "a":
+			if actualValue, expectedValue := value, 1; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case "b":
+			if actualValue, expectedValue := value, 2; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case "c":
+			if actualValue, expectedValue := value, 3; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			t.Errorf("Too many")
+		}
+		if actualValue, expectedValue := value, 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("Got %v expected %v", actualValue, expectedValue)
 	}
 }
 

diff --git a/sets/treeset/treeset.go b/sets/treeset/treeset.go
@@ -35,7 +35,7 @@ import (
 func assertInterfaceImplementation() {
 	var _ sets.Set = (*Set)(nil)
 	var _ containers.EnumerableWithIndex = (*Set)(nil)
-	var _ containers.IteratorWithIndex = (*Iterator)(nil)
+	var _ containers.ReverseIteratorWithIndex = (*Iterator)(nil)
 }
 
 // Set holds elements in a red-black tree
@@ -110,21 +110,34 @@ func (set *Set) Values() []interface{} {
 type Iterator struct {
 	index    int
 	iterator rbt.Iterator
+	tree     *rbt.Tree
 }
 
 // Iterator holding the iterator's state
 func (set *Set) Iterator() Iterator {
-	return Iterator{index: -1, iterator: set.tree.Iterator()}
+	return Iterator{index: -1, iterator: set.tree.Iterator(), tree: set.tree}
 }
 
 // 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 index and value can be retrieved by Index() and Value().
 // Modifies the state of the iterator.
 func (iterator *Iterator) Next() bool {
-	iterator.index++
+	if iterator.index < iterator.tree.Size() {
+		iterator.index++
+	}
 	return iterator.iterator.Next()
 }
 
+// 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 Index() and Value().
+// Modifies the state of the iterator.
+func (iterator *Iterator) Prev() bool {
+	if iterator.index >= 0 {
+		iterator.index--
+	}
+	return iterator.iterator.Prev()
+}
+
 // Value returns the current element's value.
 // Does not modify the state of the iterator.
 func (iterator *Iterator) Value() interface{} {

diff --git a/sets/treeset/treeset_test.go b/sets/treeset/treeset_test.go
@@ -192,10 +192,25 @@ func TestSetChaining(t *testing.T) {
 	set.Add("c", "a", "b")
 }
 
-func TestSetIterator(t *testing.T) {
+func TestSetIteratorNextOnEmpty(t *testing.T) {
 	set := NewWithStringComparator()
-	set.Add("c", "a", "b")
+	it := set.Iterator()
+	for it.Next() {
+		t.Errorf("Shouldn't iterate on empty set")
+	}
+}
 
+func TestSetIteratorPrevOnEmpty(t *testing.T) {
+	set := NewWithStringComparator()
+	it := set.Iterator()
+	for it.Prev() {
+		t.Errorf("Shouldn't iterate on empty set")
+	}
+}
+
+func TestSetIteratorNext(t *testing.T) {
+	set := NewWithStringComparator()
+	set.Add("c", "a", "b")
 	it := set.Iterator()
 	count := 0
 	for it.Next() {
@@ -225,11 +240,41 @@ func TestSetIterator(t *testing.T) {
 	if actualValue, expectedValue := count, 3; actualValue != expectedValue {
 		t.Errorf("Got %v expected %v", actualValue, expectedValue)
 	}
+}
 
-	set.Clear()
-	it = set.Iterator()
+func TestSetIteratorPrev(t *testing.T) {
+	set := NewWithStringComparator()
+	set.Add("c", "a", "b")
+	it := set.Iterator()
+	for it.Prev() {
+	}
+	count := 0
 	for it.Next() {
-		t.Errorf("Shouldn't iterate on empty set")
+		count++
+		index := it.Index()
+		value := it.Value()
+		switch index {
+		case 0:
+			if actualValue, expectedValue := value, "a"; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case 1:
+			if actualValue, expectedValue := value, "b"; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case 2:
+			if actualValue, expectedValue := value, "c"; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			t.Errorf("Too many")
+		}
+		if actualValue, expectedValue := index, count-1; actualValue != expectedValue {
+			t.Errorf("Got %v expected %v", actualValue, expectedValue)
+		}
+	}
+	if actualValue, expectedValue := count, 3; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
 	}
 }
 

diff --git a/trees/redblacktree/redblacktree_test.go b/trees/redblacktree/redblacktree_test.go
@@ -407,7 +407,7 @@ func TestRedBlackTreeIterator3Prev(t *testing.T) {
 	}
 }
 
-func TestRedBlackTreeIterator4(t *testing.T) {
+func TestRedBlackTreeIterator4Next(t *testing.T) {
 	tree := NewWithIntComparator()
 	tree.Put(13, 5)
 	tree.Put(8, 3)
@@ -419,7 +419,6 @@ func TestRedBlackTreeIterator4(t *testing.T) {
 	tree.Put(6, 2)
 	tree.Put(22, 8)
 	tree.Put(27, 10)
-
 	// │           ┌── 27
 	// │       ┌── 25
 	// │       │   └── 22
@@ -430,10 +429,7 @@ func TestRedBlackTreeIterator4(t *testing.T) {
 	//     └── 8
 	//         │   ┌── 6
 	//         └── 1
-
 	it := tree.Iterator()
-
-	// Iterator (next)
 	count := 0
 	for it.Next() {
 		count++
@@ -452,8 +448,34 @@ func TestRedBlackTreeIterator4(t *testing.T) {
 	if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue {
 		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
 	}
+}
 
-	// Iterator (prev)
+func TestRedBlackTreeIterator4(t *testing.T) {
+	tree := NewWithIntComparator()
+	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)
+	// │           ┌── 27
+	// │       ┌── 25
+	// │       │   └── 22
+	// │   ┌── 17
+	// │   │   └── 15
+	// └── 13
+	//     │   ┌── 11
+	//     └── 8
+	//         │   ┌── 6
+	//         └── 1
+	it := tree.Iterator()
+	count := tree.Size()
+	for it.Next() {
+	}
 	for it.Prev() {
 		count--
 		value := it.Value()