ℹ️ New Book: Build Your Own Database
redis/13/avl.cpp
#include "avl.h"
static uint32_t avl_depth(AVLNode *node) {
return node ? node->depth : 0;
}
static uint32_t avl_cnt(AVLNode *node) {
return node ? node->cnt : 0;
}
static uint32_t max(uint32_t lhs, uint32_t rhs) {
return lhs < rhs ? rhs : lhs;
}
// maintaining the depth and cnt field
static void avl_update(AVLNode *node) {
node->depth = 1 + max(avl_depth(node->left), avl_depth(node->right));
node->cnt = 1 + avl_cnt(node->left) + avl_cnt(node->right);
}
static AVLNode *rot_left(AVLNode *node) {
AVLNode *new_node = node->right;
if (new_node->left) {
new_node->left->parent = node;
}
node->right = new_node->left;
new_node->left = node;
new_node->parent = node->parent;
node->parent = new_node;
avl_update(node);
avl_update(new_node);
return new_node;
}
static AVLNode *rot_right(AVLNode *node) {
AVLNode *new_node = node->left;
if (new_node->right) {
new_node->right->parent = node;
}
node->left = new_node->right;
new_node->right = node;
new_node->parent = node->parent;
node->parent = new_node;
avl_update(node);
avl_update(new_node);
return new_node;
}
// the left subtree is too deep
static AVLNode *avl_fix_left(AVLNode *root) {
if (avl_depth(root->left->left) < avl_depth(root->left->right)) {
root->left = rot_left(root->left);
}
return rot_right(root);
}
// the right subtree is too deep
static AVLNode *avl_fix_right(AVLNode *root) {
if (avl_depth(root->right->right) < avl_depth(root->right->left)) {
root->right = rot_right(root->right);
}
return rot_left(root);
}
// fix imbalanced nodes and maintain invariants until the root is reached
AVLNode *avl_fix(AVLNode *node) {
while (true) {
avl_update(node);
uint32_t l = avl_depth(node->left);
uint32_t r = avl_depth(node->right);
AVLNode **from = NULL;
if (node->parent) {
from = (node->parent->left == node)
? &node->parent->left : &node->parent->right;
}
if (l == r + 2) {
node = avl_fix_left(node);
} else if (l + 2 == r) {
node = avl_fix_right(node);
}
if (!from) {
return node;
}
*from = node;
node = node->parent;
}
}
// detach a node and returns the new root of the tree
AVLNode *avl_del(AVLNode *node) {
if (node->right == NULL) {
// no right subtree, replace the node with the left subtree
// link the left subtree to the parent
AVLNode *parent = node->parent;
if (node->left) {
node->left->parent = parent;
}
if (parent) {
// attach the left subtree to the parent
(parent->left == node ? parent->left : parent->right) = node->left;
return avl_fix(parent);
} else {
// removing root?
return node->left;
}
} else {
// swap the node with its next sibling
AVLNode *victim = node->right;
while (victim->left) {
victim = victim->left;
}
AVLNode *root = avl_del(victim);
*victim = *node;
if (victim->left) {
victim->left->parent = victim;
}
if (victim->right) {
victim->right->parent = victim;
}
AVLNode *parent = node->parent;
if (parent) {
(parent->left == node ? parent->left : parent->right) = victim;
return root;
} else {
// removing root?
return victim;
}
}
}
// offset into the succeeding or preceding node.
// note: the worst-case is O(log(n)) regardless of how long the offset is.
AVLNode *avl_offset(AVLNode *node, int64_t offset) {
int64_t pos = 0; // relative to the starting node
while (offset != pos) {
if (pos < offset && pos + avl_cnt(node->right) >= offset) {
// the target is inside the right subtree
node = node->right;
pos += avl_cnt(node->left) + 1;
} else if (pos > offset && pos - avl_cnt(node->left) <= offset) {
// the target is inside the left subtree
node = node->left;
pos -= avl_cnt(node->right) + 1;
} else {
// go to the parent
AVLNode *parent = node->parent;
if (!parent) {
return NULL;
}
if (parent->right == node) {
pos -= avl_cnt(node->left) + 1;
} else {
pos += avl_cnt(node->right) + 1;
}
node = parent;
}
}
return node;
}
See also:
codecrafters.io offers “Build Your Own X” courses in many programming languages.
Including Redis, Git, SQLite, Docker, and more.
Check it out
codecrafters.io offers “Build Your Own X” courses in many programming languages.
Including Redis, Git, SQLite, Docker, and more.