diff options
| -rw-r--r-- | lib/containers/rbtree.go | 25 |
1 files changed, 23 insertions, 2 deletions
diff --git a/lib/containers/rbtree.go b/lib/containers/rbtree.go index 36da29d..ab79a26 100644 --- a/lib/containers/rbtree.go +++ b/lib/containers/rbtree.go @@ -34,8 +34,9 @@ func (node *RBNode[V]) getColor() Color { } type RBTree[K Ordered[K], V any] struct { - KeyFn func(V) K - root *RBNode[V] + KeyFn func(V) K + AttrFn func(*RBNode[V]) + root *RBNode[V] } func (t *RBTree[K, V]) Walk(fn func(*RBNode[V]) error) error { @@ -235,6 +236,16 @@ func (t *RBTree[K, V]) parentChild(node *RBNode[V]) **RBNode[V] { } } +func (t *RBTree[K, V]) updateAttr(node *RBNode[V]) { + if t.AttrFn == nil { + return + } + for node != nil { + t.AttrFn(node) + node = node.Parent + } +} + func (t *RBTree[K, V]) leftRotate(x *RBNode[V]) { // p p // | | @@ -266,6 +277,8 @@ func (t *RBTree[K, V]) leftRotate(x *RBNode[V]) { b.Parent = x } x.Right = b + + t.updateAttr(x) } func (t *RBTree[K, V]) rightRotate(y *RBNode[V]) { @@ -298,6 +311,8 @@ func (t *RBTree[K, V]) rightRotate(y *RBNode[V]) { b.Parent = y } y.Left = b + + t.updateAttr(y) } func (t *RBTree[K, V]) Insert(val V) { @@ -322,6 +337,7 @@ func (t *RBTree[K, V]) Insert(val V) { } else { parent.Right = node } + t.updateAttr(node) // Re-balance // @@ -382,6 +398,8 @@ func (t *RBTree[K, V]) Delete(key K) { // This is closely based on the algorithm presented in CLRS // 3e. + // phase 1 + var nodeToRebalance *RBNode[V] var nodeToRebalanceParent *RBNode[V] // in case 'nodeToRebalance' is nil, which it can be needsRebalance := nodeToDelete.Color == Black @@ -460,6 +478,9 @@ func (t *RBTree[K, V]) Delete(key K) { needsRebalance = next.Color == Black next.Color = nodeToDelete.Color } + t.updateAttr(nodeToRebalanceParent) + + // phase 2 if needsRebalance { node := nodeToRebalance |