mirror of
https://github.com/rdmsr/zag.git
synced 2026-10-27 13:58:20 +00:00
Loïc Grégoire
405 lines
12 KiB
Zig
405 lines
12 KiB
Zig
//! AVL tree impLementation
|
|
//! Invariant: For every node, the difference in height between its left
|
|
//! and right subtrees is strictly limited to -1, 0 or 1.
|
|
const bst = @import("bst.zig");
|
|
const std = @import("std");
|
|
|
|
// Left-heavy
|
|
const minus_one: u2 = 1;
|
|
// Balanced
|
|
const zero: u2 = 0;
|
|
// Right-heavy
|
|
const one: u2 = 2;
|
|
|
|
fn get_bf(node: *bst.Node) u2 {
|
|
return node.parent.tag();
|
|
}
|
|
|
|
fn set_bf(node: *bst.Node, bf: u2) void {
|
|
node.parent.set_tag(bf);
|
|
}
|
|
|
|
pub fn AVLTree(
|
|
comptime cmp: fn (*const bst.Node, *const bst.Node) std.math.Order,
|
|
) type {
|
|
return struct {
|
|
tree: bst.BST(cmp),
|
|
|
|
const Tree = @This();
|
|
|
|
pub fn init() Tree {
|
|
return Tree{
|
|
.tree = bst.BST(cmp).init(),
|
|
};
|
|
}
|
|
|
|
fn rotate_left(self: *Tree, x: *bst.Node, z: *bst.Node) *bst.Node {
|
|
x.right = z.left;
|
|
|
|
if (!self.tree.is_nil(z.left)) {
|
|
z.left.parent.set_ptr(x);
|
|
}
|
|
|
|
z.left = x;
|
|
|
|
x.parent.set_ptr(z);
|
|
|
|
if (get_bf(z) == zero) {
|
|
set_bf(x, one);
|
|
set_bf(z, minus_one);
|
|
} else {
|
|
set_bf(x, zero);
|
|
set_bf(z, zero);
|
|
}
|
|
|
|
return z;
|
|
}
|
|
|
|
fn rotate_right(self: *Tree, x: *bst.Node, z: *bst.Node) *bst.Node {
|
|
x.left = z.right;
|
|
|
|
if (!self.tree.is_nil(z.right)) {
|
|
z.right.parent.set_ptr(x);
|
|
}
|
|
|
|
z.right = x;
|
|
|
|
x.parent.set_ptr(z);
|
|
|
|
if (get_bf(z) == zero) {
|
|
set_bf(x, minus_one);
|
|
set_bf(z, one);
|
|
} else {
|
|
set_bf(x, zero);
|
|
set_bf(z, zero);
|
|
}
|
|
|
|
return z;
|
|
}
|
|
|
|
fn rotate_right_left(
|
|
self: *Tree,
|
|
x: *bst.Node,
|
|
z: *bst.Node,
|
|
) *bst.Node {
|
|
var y = z.left;
|
|
var t2 = y.right;
|
|
|
|
// First, rotate right at Z
|
|
z.left = t2;
|
|
if (!self.tree.is_nil(t2)) {
|
|
t2.parent.set_ptr(z);
|
|
}
|
|
|
|
y.right = z;
|
|
|
|
z.parent.set_ptr(y);
|
|
|
|
const t1 = y.left;
|
|
|
|
// second, rotate left at X
|
|
x.right = t1;
|
|
|
|
if (!self.tree.is_nil(t1)) {
|
|
t1.parent.set_ptr(x);
|
|
}
|
|
|
|
y.left = x;
|
|
x.parent.set_ptr(y);
|
|
|
|
switch (get_bf(y)) {
|
|
zero => {
|
|
set_bf(x, zero);
|
|
set_bf(z, zero);
|
|
},
|
|
|
|
one => {
|
|
set_bf(x, minus_one);
|
|
set_bf(z, zero);
|
|
},
|
|
|
|
else => {
|
|
set_bf(x, zero);
|
|
set_bf(z, one);
|
|
},
|
|
}
|
|
|
|
set_bf(y, zero);
|
|
|
|
return y;
|
|
}
|
|
|
|
fn rotate_left_right(
|
|
self: *Tree,
|
|
x: *bst.Node,
|
|
z: *bst.Node,
|
|
) *bst.Node {
|
|
var y = z.right;
|
|
var t2 = y.left;
|
|
|
|
// First, rotate left at Z
|
|
z.right = t2;
|
|
if (!self.tree.is_nil(t2)) {
|
|
t2.parent.set_ptr(z);
|
|
}
|
|
|
|
y.left = z;
|
|
|
|
z.parent.set_ptr(y);
|
|
|
|
const t1 = y.right;
|
|
|
|
// second, rotate right at X
|
|
x.left = t1;
|
|
|
|
if (!self.tree.is_nil(t1)) {
|
|
t1.parent.set_ptr(x);
|
|
}
|
|
|
|
y.right = x;
|
|
x.parent.set_ptr(y);
|
|
|
|
switch (get_bf(y)) {
|
|
zero => {
|
|
set_bf(x, zero);
|
|
set_bf(z, zero);
|
|
},
|
|
|
|
minus_one => {
|
|
set_bf(x, one);
|
|
set_bf(z, zero);
|
|
},
|
|
|
|
else => {
|
|
set_bf(x, zero);
|
|
set_bf(z, minus_one);
|
|
},
|
|
}
|
|
|
|
set_bf(y, zero);
|
|
|
|
return y;
|
|
}
|
|
|
|
/// Insert the node `elem` into the tree.
|
|
pub fn insert(self: *Tree, elem: *bst.Node) !void {
|
|
if (self.tree.is_empty()) {
|
|
self.tree.insert(elem) catch {
|
|
// This can't ever error since the tree is empty, so the
|
|
// element can't already exist.
|
|
unreachable;
|
|
};
|
|
|
|
set_bf(elem, zero);
|
|
return;
|
|
}
|
|
|
|
self.tree.insert(elem) catch |err| {
|
|
if (err != error.AlreadyExists) unreachable;
|
|
return error.AlreadyExists;
|
|
};
|
|
|
|
set_bf(elem, zero);
|
|
|
|
// Now the tree may be unbalanced, rebalance it.
|
|
var node = elem;
|
|
var parent = elem.parent.ptr();
|
|
var n = &bst.nil;
|
|
|
|
while (!self.tree.is_nil(parent)) : (parent = parent.parent.ptr()) {
|
|
const orig_parent = parent.parent.ptr();
|
|
|
|
if (node == parent.right) {
|
|
if (get_bf(parent) == one) {
|
|
// This node has been inserted on the right of the
|
|
// parent, and the parent already has a balance factor
|
|
// of 1 (right-heavy), the parent's BF would then
|
|
// become +2, which breaks the invariant.
|
|
// We need to rebalance.
|
|
|
|
if (get_bf(node) == minus_one) {
|
|
n = self.rotate_right_left(parent, node);
|
|
} else {
|
|
// This is a case where the tree looks like this:
|
|
// 1
|
|
// \
|
|
// 2
|
|
// \
|
|
// 3
|
|
// We need to rotate the subtree so it looks like:
|
|
// 2
|
|
// / \
|
|
// 1 3
|
|
// In the code, `parent` would be 1 and `node` would
|
|
// be 2
|
|
n = self.rotate_left(parent, node);
|
|
}
|
|
} else {
|
|
// The balance factor is either 0 or -1, increase it.
|
|
// If it was -1, now it is 0 so the subtree is
|
|
// perfectly balanced, we can stop.
|
|
// If it was 0, now it is +1, so we may need to check
|
|
// higher up for imbalances.
|
|
|
|
if (get_bf(parent) == minus_one) {
|
|
set_bf(parent, zero);
|
|
break;
|
|
}
|
|
|
|
set_bf(parent, one);
|
|
node = parent;
|
|
continue;
|
|
}
|
|
}
|
|
// This node is the left child.
|
|
// The logic is the same as the other case but inverted.
|
|
else {
|
|
if (get_bf(parent) == minus_one) {
|
|
// This node has been inserted on the left of the
|
|
// parent, and the parent already has a balance factor
|
|
// of -1 (left-heavy), the parent's balance factor would
|
|
// then become -2, which breaks the invariant.
|
|
// We need to rebalance.
|
|
|
|
if (get_bf(node) == one) {
|
|
n = self.rotate_left_right(parent, node);
|
|
} else {
|
|
n = self.rotate_right(parent, node);
|
|
}
|
|
} else {
|
|
if (get_bf(parent) == one) {
|
|
set_bf(parent, zero);
|
|
break;
|
|
}
|
|
|
|
set_bf(parent, minus_one);
|
|
node = parent;
|
|
continue;
|
|
}
|
|
}
|
|
|
|
// Adopt the new subtree root
|
|
n.parent.set_ptr(orig_parent);
|
|
|
|
if (!self.tree.is_nil(orig_parent)) {
|
|
if (parent == orig_parent.left) {
|
|
orig_parent.left = n;
|
|
} else {
|
|
orig_parent.right = n;
|
|
}
|
|
} else {
|
|
self.tree.root = n;
|
|
}
|
|
|
|
break;
|
|
}
|
|
}
|
|
|
|
/// Delete the node `elem` from the tree.
|
|
pub fn delete(self: *Tree, elem: *bst.Node) void {
|
|
var node = elem;
|
|
var original_parent = &bst.nil;
|
|
var parent = node.parent.ptr();
|
|
var was_on_left = !self.tree.is_nil(parent) and parent.left == node;
|
|
var b = zero;
|
|
|
|
if (self.tree.is_nil(elem.left)) {
|
|
self.tree.transplant(elem, elem.right);
|
|
} else if (self.tree.is_nil(elem.right)) {
|
|
self.tree.transplant(elem, elem.left);
|
|
} else {
|
|
var succ = self.tree.successor(node);
|
|
parent = succ.parent.ptr();
|
|
was_on_left = !self.tree.is_nil(parent) and succ == parent.left;
|
|
|
|
if (parent != elem) {
|
|
self.tree.transplant(succ, succ.right);
|
|
succ.right = elem.right;
|
|
succ.right.parent.set_ptr(succ);
|
|
} else {
|
|
// The successor is the direct child, so start rebalancing
|
|
// from it.
|
|
parent = succ;
|
|
}
|
|
|
|
self.tree.transplant(elem, succ);
|
|
succ.left = elem.left;
|
|
succ.left.parent.set_ptr(succ);
|
|
set_bf(succ, get_bf(elem));
|
|
|
|
node = succ;
|
|
}
|
|
|
|
// Rebalance the tree.
|
|
while (!self.tree.is_nil(parent)) {
|
|
original_parent = parent.parent.ptr();
|
|
|
|
if (node == parent.left or was_on_left) {
|
|
was_on_left = false;
|
|
|
|
if (get_bf(parent) == one) {
|
|
const sibling = parent.right;
|
|
b = get_bf(sibling);
|
|
|
|
if (b == minus_one) {
|
|
node = self.rotate_right_left(parent, sibling);
|
|
} else {
|
|
node = self.rotate_left(parent, sibling);
|
|
}
|
|
} else {
|
|
if (get_bf(parent) == zero) {
|
|
set_bf(parent, one);
|
|
break;
|
|
}
|
|
|
|
node = parent;
|
|
set_bf(parent, zero);
|
|
parent = original_parent;
|
|
continue;
|
|
}
|
|
} else {
|
|
if (get_bf(parent) == minus_one) {
|
|
const sibling = parent.left;
|
|
b = get_bf(sibling);
|
|
|
|
if (b == one) {
|
|
node = self.rotate_left_right(parent, sibling);
|
|
} else {
|
|
node = self.rotate_right(parent, sibling);
|
|
}
|
|
} else {
|
|
if (get_bf(parent) == zero) {
|
|
set_bf(parent, minus_one);
|
|
break;
|
|
}
|
|
|
|
node = parent;
|
|
set_bf(node, zero);
|
|
parent = original_parent;
|
|
continue;
|
|
}
|
|
}
|
|
|
|
// Adopt the new subtree root
|
|
node.parent.set_ptr(original_parent);
|
|
|
|
if (!self.tree.is_nil(original_parent)) {
|
|
if (parent == original_parent.left) {
|
|
original_parent.left = node;
|
|
} else {
|
|
original_parent.right = node;
|
|
}
|
|
} else {
|
|
self.tree.root = node;
|
|
}
|
|
|
|
if (b == zero) {
|
|
break;
|
|
}
|
|
|
|
parent = original_parent;
|
|
}
|
|
}
|
|
};
|
|
}
|