From 69d218401163c13c474abc1d08b61fdfae63621b Mon Sep 17 00:00:00 2001 From: Emir Pasic Date: Tue, 12 Apr 2022 03:33:59 +0200 Subject: [PATCH] Implement GetNode() for BTree AVLTree and RedBlackTree and Node() function for their iterators --- trees/avltree/avltree.go | 30 ++++++++++++++-- trees/avltree/avltree_test.go | 44 +++++++++++++++++++++++ trees/avltree/iterator.go | 6 ++++ trees/btree/btree.go | 20 +++++++++++ trees/btree/btree_test.go | 46 +++++++++++++++++++++++++ trees/btree/iterator.go | 6 ++++ trees/redblacktree/iterator.go | 6 ++++ trees/redblacktree/redblacktree.go | 22 ++++++++++++ trees/redblacktree/redblacktree_test.go | 46 +++++++++++++++++++++++++ 9 files changed, 224 insertions(+), 2 deletions(-) diff --git a/trees/avltree/avltree.go b/trees/avltree/avltree.go index 1efdae4c..fd1a4053 100644 --- a/trees/avltree/avltree.go +++ b/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) diff --git a/trees/avltree/avltree_test.go b/trees/avltree/avltree_test.go index 16bf3017..9cc300a4 100644 --- a/trees/avltree/avltree_test.go +++ b/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") diff --git a/trees/avltree/iterator.go b/trees/avltree/iterator.go index d4a1585e..c4d69f4d 100644 --- a/trees/avltree/iterator.go +++ b/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() { diff --git a/trees/btree/btree.go b/trees/btree/btree.go index 5f866997..523358ad 100644 --- a/trees/btree/btree.go +++ b/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) diff --git a/trees/btree/btree_test.go b/trees/btree/btree_test.go index f4edaf1f..5bf5cb4d 100644 --- a/trees/btree/btree_test.go +++ b/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) diff --git a/trees/btree/iterator.go b/trees/btree/iterator.go index ed0cb154..c74b6464 100644 --- a/trees/btree/iterator.go +++ b/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() { diff --git a/trees/redblacktree/iterator.go b/trees/redblacktree/iterator.go index 881b67dd..79d6ab3a 100644 --- a/trees/redblacktree/iterator.go +++ b/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() { diff --git a/trees/redblacktree/redblacktree.go b/trees/redblacktree/redblacktree.go index 8301ffd7..cce9fe0a 100644 --- a/trees/redblacktree/redblacktree.go +++ b/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) diff --git a/trees/redblacktree/redblacktree_test.go b/trees/redblacktree/redblacktree_test.go index 730a1f3a..3abe6c1a 100644 --- a/trees/redblacktree/redblacktree_test.go +++ b/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")