Implement GetNode() for BTree AVLTree and RedBlackTree and Node() function for their iterators

trees/avltree/avltree.go

@@ -60,19 +60,29 @@ func (t *Tree) Put(key interface{}, value interface{}) {
 // Second return parameter is true if key was found, otherwise false.
 // Key should adhere to the comparator's type assertion, otherwise method panics.
 func (t *Tree) Get(key interface{}) (value interface{}, found bool) {
+	n := t.GetNode(key)
+	if n != nil {
+		return n.Value, true
+	}
+	return nil, false
+}
+
+// GetNode searches the node in the tree by key and returns its node or nil if key is not found in tree.
+// Key should adhere to the comparator's type assertion, otherwise method panics.
+func (t *Tree) GetNode(key interface{}) *Node {
 	n := t.Root
 	for n != nil {
 		cmp := t.Comparator(key, n.Key)
 		switch {
 		case cmp == 0:
-			return n.Value, true
+			return n
 		case cmp < 0:
 			n = n.Children[0]
 		case cmp > 0:
 			n = n.Children[1]
 		}
 	}
-	return nil, false
+	return n
 }
 
 // Remove remove the node from the tree by key.
@@ -91,6 +101,22 @@ func (t *Tree) Size() int {
 	return t.size
 }
 
+// Size returns the number of elements stored in the subtree.
+// Computed dynamically on each call, i.e. the subtree is traversed to count the number of the nodes.
+func (n *Node) Size() int {
+	if n == nil {
+		return 0
+	}
+	size := 1
+	if n.Children[0] != nil {
+		size += n.Children[0].Size()
+	}
+	if n.Children[1] != nil {
+		size += n.Children[1].Size()
+	}
+	return size
+}
+
 // Keys returns all keys in-order
 func (t *Tree) Keys() []interface{} {
 	keys := make([]interface{}, t.size)

trees/avltree/avltree_test.go

@@ -9,6 +9,50 @@ import (
 	"testing"
 )
 
+func TestAVLTreeGet(t *testing.T) {
+	tree := NewWithIntComparator()
+
+	if actualValue := tree.Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+
+	if actualValue := tree.GetNode(2).Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+
+	tree.Put(1, "x") // 1->x
+	tree.Put(2, "b") // 1->x, 2->b (in order)
+	tree.Put(1, "a") // 1->a, 2->b (in order, replacement)
+	tree.Put(3, "c") // 1->a, 2->b, 3->c (in order)
+	tree.Put(4, "d") // 1->a, 2->b, 3->c, 4->d (in order)
+	tree.Put(5, "e") // 1->a, 2->b, 3->c, 4->d, 5->e (in order)
+	tree.Put(6, "f") // 1->a, 2->b, 3->c, 4->d, 5->e, 6->f (in order)
+	//
+	//  AVLTree
+	//  │       ┌── 6
+	//  │   ┌── 5
+	//  └── 4
+	//      │   ┌── 3
+	//      └── 2
+	//          └── 1
+
+	if actualValue := tree.Size(); actualValue != 6 {
+		t.Errorf("Got %v expected %v", actualValue, 6)
+	}
+
+	if actualValue := tree.GetNode(2).Size(); actualValue != 3 {
+		t.Errorf("Got %v expected %v", actualValue, 3)
+	}
+
+	if actualValue := tree.GetNode(4).Size(); actualValue != 6 {
+		t.Errorf("Got %v expected %v", actualValue, 6)
+	}
+
+	if actualValue := tree.GetNode(7).Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+}
+
 func TestAVLTreePut(t *testing.T) {
 	tree := NewWithIntComparator()
 	tree.Put(5, "e")

trees/avltree/iterator.go

@@ -86,6 +86,12 @@ func (iterator *Iterator) Key() interface{} {
 	return iterator.node.Key
 }
 
+// Node returns the current element's node.
+// Does not modify the state of the iterator.
+func (iterator *Iterator) Node() *Node {
+	return iterator.node
+}
+
 // Begin resets the iterator to its initial state (one-before-first)
 // Call Next() to fetch the first element if any.
 func (iterator *Iterator) Begin() {

trees/btree/btree.go

@@ -95,6 +95,13 @@ func (tree *Tree) Get(key interface{}) (value interface{}, found bool) {
 	return nil, false
 }
 
+// GetNode searches the node in the tree by key and returns its node or nil if key is not found in tree.
+// Key should adhere to the comparator's type assertion, otherwise method panics.
+func (tree *Tree) GetNode(key interface{}) *Node {
+	node, _, _ := tree.searchRecursively(tree.Root, key)
+	return node
+}
+
 // Remove remove the node from the tree by key.
 // Key should adhere to the comparator's type assertion, otherwise method panics.
 func (tree *Tree) Remove(key interface{}) {
@@ -115,6 +122,19 @@ func (tree *Tree) Size() int {
 	return tree.size
 }
 
+// Size returns the number of elements stored in the subtree.
+// Computed dynamically on each call, i.e. the subtree is traversed to count the number of the nodes.
+func (node *Node) Size() int {
+	if node == nil {
+		return 0
+	}
+	size := 1
+	for _, child := range node.Children {
+		size += child.Size()
+	}
+	return size
+}
+
 // Keys returns all keys in-order
 func (tree *Tree) Keys() []interface{} {
 	keys := make([]interface{}, tree.size)

trees/btree/btree_test.go

@@ -74,6 +74,52 @@ func TestBTreeGet2(t *testing.T) {
 	}
 }
 
+func TestBTreeGet3(t *testing.T) {
+	tree := NewWithIntComparator(3)
+
+	if actualValue := tree.Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+
+	if actualValue := tree.GetNode(2).Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+
+	tree.Put(1, "x") // 1->x
+	tree.Put(2, "b") // 1->x, 2->b (in order)
+	tree.Put(1, "a") // 1->a, 2->b (in order, replacement)
+	tree.Put(3, "c") // 1->a, 2->b, 3->c (in order)
+	tree.Put(4, "d") // 1->a, 2->b, 3->c, 4->d (in order)
+	tree.Put(5, "e") // 1->a, 2->b, 3->c, 4->d, 5->e (in order)
+	tree.Put(6, "f") // 1->a, 2->b, 3->c, 4->d, 5->e, 6->f (in order)
+	tree.Put(7, "g") // 1->a, 2->b, 3->c, 4->d, 5->e, 6->f, 7->g (in order)
+
+	// BTree
+	//         1
+	//     2
+	//         3
+	// 4
+	//         5
+	//     6
+	//         7
+
+	if actualValue := tree.Size(); actualValue != 7 {
+		t.Errorf("Got %v expected %v", actualValue, 7)
+	}
+
+	if actualValue := tree.GetNode(2).Size(); actualValue != 3 {
+		t.Errorf("Got %v expected %v", actualValue, 3)
+	}
+
+	if actualValue := tree.GetNode(4).Size(); actualValue != 7 {
+		t.Errorf("Got %v expected %v", actualValue, 7)
+	}
+
+	if actualValue := tree.GetNode(8).Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+}
+
 func TestBTreePut1(t *testing.T) {
 	// https://upload.wikimedia.org/wikipedia/commons/3/33/B_tree_insertion_example.png
 	tree := NewWithIntComparator(3)

trees/btree/iterator.go

@@ -160,6 +160,12 @@ func (iterator *Iterator) Key() interface{} {
 	return iterator.entry.Key
 }
 
+// Node returns the current element's node.
+// Does not modify the state of the iterator.
+func (iterator *Iterator) Node() *Node {
+	return iterator.node
+}
+
 // Begin resets the iterator to its initial state (one-before-first)
 // Call Next() to fetch the first element if any.
 func (iterator *Iterator) Begin() {

trees/redblacktree/iterator.go

@@ -130,6 +130,12 @@ func (iterator *Iterator) Key() interface{} {
 	return iterator.node.Key
 }
 
+// Node returns the current element's node.
+// Does not modify the state of the iterator.
+func (iterator *Iterator) Node() *Node {
+	return iterator.node
+}
+
 // Begin resets the iterator to its initial state (one-before-first)
 // Call Next() to fetch the first element if any.
 func (iterator *Iterator) Begin() {

trees/redblacktree/redblacktree.go

@@ -113,6 +113,12 @@ func (tree *Tree) Get(key interface{}) (value interface{}, found bool) {
 	return nil, false
 }
 
+// GetNode searches the node in the tree by key and returns its node or nil if key is not found in tree.
+// Key should adhere to the comparator's type assertion, otherwise method panics.
+func (tree *Tree) GetNode(key interface{}) *Node {
+	return tree.lookup(key)
+}
+
 // Remove remove the node from the tree by key.
 // Key should adhere to the comparator's type assertion, otherwise method panics.
 func (tree *Tree) Remove(key interface{}) {
@@ -155,6 +161,22 @@ func (tree *Tree) Size() int {
 	return tree.size
 }
 
+// Size returns the number of elements stored in the subtree.
+// Computed dynamically on each call, i.e. the subtree is traversed to count the number of the nodes.
+func (node *Node) Size() int {
+	if node == nil {
+		return 0
+	}
+	size := 1
+	if node.Left != nil {
+		size += node.Left.Size()
+	}
+	if node.Right != nil {
+		size += node.Right.Size()
+	}
+	return size
+}
+
 // Keys returns all keys in-order
 func (tree *Tree) Keys() []interface{} {
 	keys := make([]interface{}, tree.size)

trees/redblacktree/redblacktree_test.go

@@ -10,6 +10,52 @@ import (
 	"testing"
 )
 
+func TestRedBlackTreeGet(t *testing.T) {
+	tree := NewWithIntComparator()
+
+	if actualValue := tree.Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+
+	if actualValue := tree.GetNode(2).Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+
+	tree.Put(1, "x") // 1->x
+	tree.Put(2, "b") // 1->x, 2->b (in order)
+	tree.Put(1, "a") // 1->a, 2->b (in order, replacement)
+	tree.Put(3, "c") // 1->a, 2->b, 3->c (in order)
+	tree.Put(4, "d") // 1->a, 2->b, 3->c, 4->d (in order)
+	tree.Put(5, "e") // 1->a, 2->b, 3->c, 4->d, 5->e (in order)
+	tree.Put(6, "f") // 1->a, 2->b, 3->c, 4->d, 5->e, 6->f (in order)
+
+	fmt.Println(tree)
+	//
+	//  RedBlackTree
+	//  │           ┌── 6
+	//  │       ┌── 5
+	//  │   ┌── 4
+	//  │   │   └── 3
+	//  └── 2
+	//       └── 1
+
+	if actualValue := tree.Size(); actualValue != 6 {
+		t.Errorf("Got %v expected %v", actualValue, 6)
+	}
+
+	if actualValue := tree.GetNode(4).Size(); actualValue != 4 {
+		t.Errorf("Got %v expected %v", actualValue, 4)
+	}
+
+	if actualValue := tree.GetNode(2).Size(); actualValue != 6 {
+		t.Errorf("Got %v expected %v", actualValue, 6)
+	}
+
+	if actualValue := tree.GetNode(8).Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+}
+
 func TestRedBlackTreePut(t *testing.T) {
 	tree := NewWithIntComparator()
 	tree.Put(5, "e")