From 84b2f395c68f52f62a7782ca9ad2093c93dd30a6 Mon Sep 17 00:00:00 2001 From: Devansh Singh Date: Sun, 18 Feb 2024 17:01:24 +0530 Subject: [PATCH] Add deletion method for RB Tree Signed-off-by: Devansh Singh --- set/set_test.go | 2 +- zset/rbtree.go | 139 ++++++++++++++++++++++++++++++++++++++++++-- zset/rbtree_test.go | 53 +++++++++++++++++ 3 files changed, 189 insertions(+), 5 deletions(-) create mode 100644 zset/rbtree_test.go diff --git a/set/set_test.go b/set/set_test.go index 8881a7d..fbd4ea0 100644 --- a/set/set_test.go +++ b/set/set_test.go @@ -2,7 +2,7 @@ package set import "testing" -var members = []string{"hello", "world"} +var members = []string{"hello", "world", "how", "are", "you"} func createTestSet() Set { set := NewSet() diff --git a/zset/rbtree.go b/zset/rbtree.go index acc143a..680e53f 100644 --- a/zset/rbtree.go +++ b/zset/rbtree.go @@ -82,11 +82,11 @@ func (t *RBTree) rightRotate(node *Node) { node.Parent = lnode } -func (t *RBTree) search(value string) bool { +func (t *RBTree) search(value string) (*Node, bool) { temp := t.Root for temp != nil { if temp.Value == value { - return true + return temp, true } if temp.Value > value { temp = temp.Left @@ -94,7 +94,7 @@ func (t *RBTree) search(value string) bool { temp = temp.Right } } - return false + return nil, false } func (t *RBTree) fixInsert(node *Node) { @@ -139,7 +139,7 @@ func (t *RBTree) fixInsert(node *Node) { } func (t *RBTree) insert(value string) { - if ok := t.search(value); ok { + if _, ok := t.search(value); ok { return } node := NewNode(value) @@ -168,6 +168,137 @@ func (t *RBTree) insert(value string) { t.fixInsert(node) } +func (t *RBTree) min(node *Node) *Node { + if node == nil { + return nil + } + + for node.Left != nil { + node = node.Left + } + return node +} + +func (t *RBTree) max(node *Node) *Node { + if node == nil { + return nil + } + + for node.Right != nil { + node = node.Right + } + return node +} + +func (t *RBTree) successor(node *Node) *Node { + if node == nil { + return nil + } + + if node.Right != nil { + return t.min(node.Right) + } + + successor := node.Parent + for successor != nil && node == successor.Right { + node = successor + successor = successor.Parent + } + return successor +} + +func (t *RBTree) fixDelete(node *Node) { + for node != t.Root && node.Color == BLACK { + if node == node.Parent.Left { + cousin := node.Parent.Right + if cousin.Color == RED { + cousin.Color = BLACK + node.Parent.Color = RED + t.leftRotate(node.Parent) + cousin = node.Parent.Right + } + if cousin.Left.Color == BLACK && cousin.Right.Color == BLACK { + cousin.Color = RED + node = node.Parent + } else { + if cousin.Right.Color == BLACK { + cousin.Left.Color = BLACK + cousin.Color = RED + t.rightRotate(cousin) + cousin = node.Parent.Right + } + cousin.Color = node.Parent.Color + node.Parent.Color = BLACK + cousin.Right.Color = BLACK + t.leftRotate(node.Parent) + node = t.Root + } + } else { + cousin := node.Parent.Left + if cousin.Color == RED { + cousin.Color = BLACK + node.Parent.Color = RED + t.rightRotate(node.Parent) + cousin = node.Parent.Left + } + if cousin.Left.Color == BLACK && cousin.Right.Color == BLACK { + cousin.Color = RED + node = node.Parent + } else { + if cousin.Left.Color == BLACK { + cousin.Right.Color = BLACK + cousin.Color = RED + t.leftRotate(cousin) + cousin = node.Parent.Left + } + cousin.Color = node.Parent.Color + node.Parent.Color = BLACK + node.Left.Color = BLACK + t.rightRotate(node.Parent) + node = t.Root + } + } + } + node.Color = BLACK +} + +func (t *RBTree) delete(value string) { + node, ok := t.search(value) + if !ok { + return + } + + var temp *Node + if node.Left == nil || node.Right == nil { + temp = node + } else { + temp = t.successor(node) + } + var tchild *Node + if temp.Left != nil { + tchild = temp.Left + } else { + tchild = temp.Right + } + + tchild.Parent = temp.Parent + if temp.Parent == nil { + t.Root = tchild + } else if temp == temp.Parent.Left { + temp.Parent.Left = tchild + } else { + temp.Parent.Right = tchild + } + + if temp != node { + node.Value = temp.Value + } + if temp.Color == BLACK { + t.fixDelete(tchild) + } + t.Count-- +} + func inorder(node *Node, elements chan string) { if node == nil { return diff --git a/zset/rbtree_test.go b/zset/rbtree_test.go new file mode 100644 index 0000000..aa34dd6 --- /dev/null +++ b/zset/rbtree_test.go @@ -0,0 +1,53 @@ +package zset + +import ( + "reflect" + "testing" +) + +var members = []string{"hello", "world", "how", "are", "you"} + +func createTestTree() *RBTree { + tree := NewRBTree() + for _, member := range members { + tree.insert(member) + } + return tree +} + +func TestInsert(t *testing.T) { + tree := createTestTree() + + got := tree.members() + want := []string{"are", "hello", "how", "world", "you"} + if !reflect.DeepEqual(got, want) { + t.Errorf("got %q, wanted %q", got, want) + } +} + +func TestDelete(t *testing.T) { + tree := createTestTree() + tree.delete("hello") + + got := tree.members() + want := []string{"are", "how", "world", "you"} + if !reflect.DeepEqual(got, want) { + t.Errorf("got %q, wanted %q", got, want) + } +} + +func TestSearch(t *testing.T) { + tree := createTestTree() + + _, got := tree.search("hello") + want := true + if got != want { + t.Errorf("got %t, wanted %t", got, want) + } + + _, got = tree.search("secctan") + want = false + if got != want { + t.Errorf("got %t, wanted %t", got, want) + } +}