summaryrefslogtreecommitdiff
path: root/lib/rbtree/rbtree.go
diff options
context:
space:
mode:
Diffstat (limited to 'lib/rbtree/rbtree.go')
-rw-r--r--lib/rbtree/rbtree.go477
1 files changed, 477 insertions, 0 deletions
diff --git a/lib/rbtree/rbtree.go b/lib/rbtree/rbtree.go
new file mode 100644
index 0000000..7927307
--- /dev/null
+++ b/lib/rbtree/rbtree.go
@@ -0,0 +1,477 @@
+package rbtree
+
+import (
+ "fmt"
+
+ "golang.org/x/exp/constraints"
+)
+
+type Color bool
+
+const (
+ Black = Color(false)
+ Red = Color(true)
+)
+
+type Node[V any] struct {
+ Parent, Left, Right *Node[V]
+
+ Color Color
+
+ Value V
+}
+
+func (node *Node[V]) getColor() Color {
+ if node == nil {
+ return Black
+ }
+ return node.Color
+}
+
+type Tree[K constraints.Ordered, V any] struct {
+ KeyFn func(V) K
+ root *Node[V]
+}
+
+func (t *Tree[K, V]) Walk(fn func(*Node[V]) error) error {
+ return t.root.walk(fn)
+}
+
+func (node *Node[V]) walk(fn func(*Node[V]) error) error {
+ if node == nil {
+ return nil
+ }
+ if err := node.Left.walk(fn); err != nil {
+ return err
+ }
+ if err := fn(node); err != nil {
+ return err
+ }
+ if err := node.Right.walk(fn); err != nil {
+ return err
+ }
+ return nil
+}
+
+// Search the tree for a value that satisfied the given callbackk
+// function. A return value of 0 means to to return this value; <0
+// means to go left on the tree (the value is too high), >0 means to
+// go right on th etree (the value is too low).
+//
+// +-----+
+// | v=8 | == 0 : this is it
+// +-----+
+// / \
+// / \
+// <0 : go left >0 : go right
+// / \
+// +---+ +---+
+// | 7 | | 9 |
+// +---+ +---+
+//
+// Returns nil if no such value is found.
+//
+// Search is good for advanced lookup, like when a range of values is
+// acceptable. For simple exact-value lookup, use Lookup.
+func (t *Tree[K, V]) Search(fn func(V) int) *Node[V] {
+ ret, _ := t.root.search(fn)
+ return ret
+}
+
+func (node *Node[V]) search(fn func(V) int) (exact, nearest *Node[V]) {
+ var prev *Node[V]
+ for {
+ if node == nil {
+ return nil, prev
+ }
+ direction := fn(node.Value)
+ prev = node
+ switch {
+ case direction < 0:
+ node = node.Left
+ case direction == 0:
+ return node, nil
+ case direction > 0:
+ node = node.Right
+ }
+ }
+}
+
+func (t *Tree[K, V]) exactKey(key K) func(V) int {
+ return func(val V) int {
+ valKey := t.KeyFn(val)
+ switch {
+ case key < valKey:
+ return -1
+ case key > valKey:
+ return 1
+ default: // key == valKey:
+ return 0
+ }
+ }
+}
+
+// Lookup looks up the value for an exact key. If no such value
+// exists, nil is returned.
+func (t *Tree[K, V]) Lookup(key K) *Node[V] {
+ return t.Search(t.exactKey(key))
+}
+
+// Min returns the minimum value stored in the tree, or nil if the
+// tree is empty.
+func (t *Tree[K, V]) Min() *Node[V] {
+ return t.root.min()
+}
+
+func (node *Node[V]) min() *Node[V] {
+ if node == nil {
+ return nil
+ }
+ for {
+ if node.Left == nil {
+ return node
+ }
+ node = node.Left
+ }
+}
+
+// Max returns the maximum value stored in the tree, or nil if the
+// tree is empty.
+func (t *Tree[K, V]) Max() *Node[V] {
+ return t.root.max()
+}
+
+func (node *Node[V]) max() *Node[V] {
+ if node == nil {
+ return nil
+ }
+ for {
+ if node.Right == nil {
+ return node
+ }
+ node = node.Right
+ }
+}
+
+func (t *Tree[K, V]) Next(cur *Node[V]) *Node[V] {
+ return cur.next()
+}
+
+func (cur *Node[V]) next() *Node[V] {
+ if cur.Right != nil {
+ return cur.Right.min()
+ }
+ child, parent := cur, cur.Parent
+ for parent != nil && child == parent.Right {
+ child, parent = parent, parent.Parent
+ }
+ return parent
+}
+
+func (t *Tree[K, V]) Prev(cur *Node[V]) *Node[V] {
+ return cur.prev()
+}
+
+func (cur *Node[V]) prev() *Node[V] {
+ if cur.Left != nil {
+ return cur.Left.max()
+ }
+ child, parent := cur, cur.Parent
+ for parent != nil && child == parent.Left {
+ child, parent = parent, parent.Parent
+ }
+ return parent
+}
+
+func (t *Tree[K, V]) parentChild(node *Node[V]) **Node[V] {
+ switch {
+ case node.Parent == nil:
+ return &t.root
+ case node.Parent.Left == node:
+ return &node.Parent.Left
+ case node.Parent.Right == node:
+ return &node.Parent.Right
+ default:
+ panic(fmt.Errorf("node %p is not a child of its parent %p", node, node.Parent))
+ }
+}
+
+func (t *Tree[K, V]) leftRotate(x *Node[V]) {
+ // p p
+ // | |
+ // +---+ +---+
+ // | x | | y |
+ // +---+ +---+
+ // / \ => / \
+ // a +---+ +---+ c
+ // | y | | x |
+ // +---+ +---+
+ // / \ / \
+ // b c a b
+
+ // Define 'p', 'x', 'y', and 'b' per the above diagram.
+ p := x.Parent
+ pChild := t.parentChild(x)
+ y := x.Right
+ b := y.Left
+
+ // Move things around
+
+ y.Parent = p
+ *pChild = y
+
+ x.Parent = y
+ y.Left = x
+
+ if b != nil {
+ b.Parent = x
+ }
+ x.Right = b
+}
+
+func (t *Tree[K, V]) rightRotate(y *Node[V]) {
+ // | |
+ // +---+ +---+
+ // | y | | x |
+ // +---+ +---+
+ // / \ => / \
+ // +---+ c a +---+
+ // | x | | y |
+ // +---+ +---+
+ // / \ / \
+ // a b b c
+
+ // Define 'p', 'x', 'y', and 'b' per the above diagram.
+ p := y.Parent
+ pChild := t.parentChild(y)
+ x := y.Left
+ b := x.Right
+
+ // Move things around
+
+ x.Parent = p
+ *pChild = x
+
+ y.Parent = x
+ x.Right = y
+
+ if b != nil {
+ b.Parent = y
+ }
+ y.Left = b
+}
+
+func (t *Tree[K, V]) Insert(val V) {
+ // Naive-insert
+
+ key := t.KeyFn(val)
+ exact, parent := t.root.search(t.exactKey(key))
+ if exact != nil {
+ exact.Value = val
+ return
+ }
+
+ node := &Node[V]{
+ Color: Red,
+ Parent: parent,
+ Value: val,
+ }
+ if parent == nil {
+ t.root = node
+ } else if key < t.KeyFn(parent.Value) {
+ parent.Left = node
+ } else {
+ parent.Right = node
+ }
+
+ // Re-balance
+
+ for node.Parent.getColor() == Red {
+ if node.Parent == node.Parent.Parent.Left {
+ uncle := node.Parent.Parent.Right
+ if uncle.getColor() == Red {
+ node.Parent.Color = Black
+ uncle.Color = Black
+ node.Parent.Parent.Color = Red
+ node = node.Parent.Parent
+ } else {
+ if node == node.Parent.Right {
+ node = node.Parent
+ t.leftRotate(node)
+ }
+ node.Parent.Color = Black
+ node.Parent.Parent.Color = Red
+ t.rightRotate(node.Parent.Parent)
+ }
+ } else {
+ uncle := node.Parent.Parent.Left
+ if uncle.getColor() == Red {
+ node.Parent.Color = Black
+ uncle.Color = Black
+ node.Parent.Parent.Color = Red
+ node = node.Parent.Parent
+ } else {
+ if node == node.Parent.Left {
+ node = node.Parent
+ t.rightRotate(node)
+ }
+ node.Parent.Color = Black
+ node.Parent.Parent.Color = Red
+ t.leftRotate(node.Parent.Parent)
+ }
+ }
+ }
+ t.root.Color = Black
+}
+
+func (t *Tree[K, V]) transplant(old, new *Node[V]) {
+ *t.parentChild(old) = new
+ if new != nil {
+ new.Parent = old.Parent
+ }
+}
+
+func (t *Tree[K, V]) Delete(key K) {
+ nodeToDelete := t.Lookup(key)
+ if nodeToDelete == nil {
+ return
+ }
+
+ var nodeToRebalance *Node[V]
+ var nodeToRebalanceParent *Node[V] // in case 'nodeToRebalance' is nil, which it can be
+ needsRebalance := nodeToDelete.Color == Black
+
+ switch {
+ case nodeToDelete.Left == nil:
+ nodeToRebalance = nodeToDelete.Right
+ nodeToRebalanceParent = nodeToDelete.Parent
+ t.transplant(nodeToDelete, nodeToDelete.Right)
+ case nodeToDelete.Right == nil:
+ nodeToRebalance = nodeToDelete.Left
+ nodeToRebalanceParent = nodeToDelete.Parent
+ t.transplant(nodeToDelete, nodeToDelete.Left)
+ default:
+ // The node being deleted has a child on both sides,
+ // so we've go to reshuffle the parents a bit to make
+ // room for those children.
+ next := nodeToDelete.next()
+ if next.Parent == nodeToDelete {
+ // p p
+ // | |
+ // +-----+ +-----+
+ // | ntd | | nxt |
+ // +-----+ +-----+
+ // / \ => / \
+ // a +-----+ a b
+ // | nxt |
+ // +-----+
+ // / \
+ // nil b
+ nodeToRebalance = next.Right
+ nodeToRebalanceParent = next
+
+ *t.parentChild(nodeToDelete) = next
+ next.Parent = nodeToDelete.Parent
+
+ next.Left = nodeToDelete.Left
+ next.Left.Parent = next
+ } else {
+ // p p
+ // | |
+ // +-----+ +-----+
+ // | ntd | | nxt |
+ // +-----+ +-----+
+ // / \ / \
+ // a x a x
+ // / \ => / \
+ // y z y z
+ // / \ / \
+ // +-----+ c b c
+ // | nxt |
+ // +-----+
+ // / \
+ // nil b
+ y := next.Parent
+ b := next.Right
+ nodeToRebalance = b
+ nodeToRebalanceParent = y
+
+ *t.parentChild(nodeToDelete) = next
+ next.Parent = nodeToDelete.Parent
+
+ next.Left = nodeToDelete.Left
+ next.Left.Parent = next
+
+ next.Right = nodeToDelete.Right
+ next.Right.Parent = next
+
+ y.Left = b
+ if b != nil {
+ b.Parent = y
+ }
+ }
+
+ // idk
+ needsRebalance = next.Color == Black
+ next.Color = nodeToDelete.Color
+ }
+
+ if needsRebalance {
+ node := nodeToRebalance
+ nodeParent := nodeToRebalanceParent
+ for node != t.root && node.getColor() == Black {
+ if node == nodeParent.Left {
+ sibling := nodeParent.Right
+ if sibling.getColor() == Red {
+ sibling.Color = Black
+ nodeParent.Color = Red
+ t.leftRotate(nodeParent)
+ sibling = nodeParent.Right
+ }
+ if sibling.Left.getColor() == Black && sibling.Right.getColor() == Black {
+ sibling.Color = Red
+ node, nodeParent = nodeParent, nodeParent.Parent
+ } else {
+ if sibling.Right.getColor() == Black {
+ sibling.Left.Color = Black
+ sibling.Color = Red
+ t.rightRotate(sibling)
+ sibling = nodeParent.Right
+ }
+ sibling.Color = nodeParent.Color
+ nodeParent.Color = Black
+ sibling.Right.Color = Black
+ t.leftRotate(nodeParent)
+ node, nodeParent = t.root, nil
+ }
+ } else {
+ sibling := nodeParent.Left
+ if sibling.getColor() == Red {
+ sibling.Color = Black
+ nodeParent.Color = Red
+ t.rightRotate(nodeParent)
+ sibling = nodeParent.Left
+ }
+ if sibling.Right.getColor() == Black && sibling.Left.getColor() == Black {
+ sibling.Color = Red
+ node, nodeParent = nodeParent, nodeParent.Parent
+ } else {
+ if sibling.Left.getColor() == Black {
+ sibling.Right.Color = Black
+ sibling.Color = Red
+ t.leftRotate(sibling)
+ sibling = nodeParent.Left
+ }
+ sibling.Color = nodeParent.Color
+ nodeParent.Color = Black
+ sibling.Left.Color = Black
+ t.rightRotate(nodeParent)
+ node, nodeParent = t.root, nil
+ }
+ }
+ }
+ if node != nil {
+ node.Color = Black
+ }
+ }
+}
generated by cgit v1.2.3-54-g00ecf (git 2.45.2) at 2024-07-24 13:53:38 +0000