I never remember without looking how to write RB or AVL balancing. but I wanted to make sure I would be able to still write a reasonable dynamically balanced binary tree. If it is asked of me in interviews.
So here is my attempt at it.
#include <cassert>
#include <iostream>
#include <vector>
template<typename K, typename V>
class BalancedBinaryTree {
struct Node {
K key;
V val;
int depth;
Node* left;
Node* right;
};
Node* root = nullptr;
Node* try_lookup(Node* n, K key) {
if (!n || n->key == key)
return n;
if (key > n->key)
return try_lookup(n->right, key);
return try_lookup(n->left, key);
}
struct Depth2 {
int left;
int right;
int max() { return std::max(left, right); }
};
int depth1(Node* n) {
if (!n)
return 0;
return n->depth;
}
Depth2 depth2(Node* n) {
if (!n)
return {0, 0};
return {1 + depth1(n->left), 1 + depth1(n->right)};
}
/*
A
/ \
B C
/ \
D E
to
C
/ \
A E
/ \
B D
*/
void rotate_left(Node **at) {
Node* n = *at;
Node *new_root = n->left;
n->left = n->left->right;
new_root->right = n;
n->depth = depth2(n).max();
new_root->depth = depth2(new_root).max();
*at = new_root;
}
// The other side
void rotate_right(Node **at) {
Node *n = *at;
Node *new_root = n->right;
n->right = n->right->left;
new_root->left = n;
n->depth = depth2(n).max();
new_root->depth = depth2(new_root).max();
*at = new_root;
}
void local_rebalance(Node** at) {
Node* n = *at;
Depth2 left = depth2(n->left);
Depth2 right = depth2(n->right);
int diff = left.max() - right.max();
if (diff >= -1 && diff <= 1) {
Depth2 res = {1 + left.max(), 1 + right.max()};
n->depth = res.max();
return;
}
if (left.max() > right.max())
return rotate_left(at);
else
return rotate_right(at);
}
template <typename FuncT> void update_impl(Node **at, K key, FuncT func) {
Node *n = *at;
if (!n || n->key == key)
return func(at);
if (key > n->key)
update_impl(&n->right, key, func);
else
update_impl(&n->left, key, func);
local_rebalance(at);
}
void release_from_tree(Node **at) {
Node *n = *at;
auto remove_and_replace = [&](Node* by) {
delete n;
*at = by;
return;
};
if (!n->right)
return remove_and_replace(n->left);
if (!n->left)
return remove_and_replace(n->right);
rotate_left(at);
Node *new_root = *at;
release_from_tree(&new_root->left);
local_rebalance(at);
}
public:
bool insert(K key, V val) {
bool res = false;
update_impl(&root, key, [&](Node **at) {
assert(at);
Node *n = *at;
if (n) {
res = false;
n->val = std::move(val);
return;
}
res = true;
*at = new Node{key, std::move(val), 1, nullptr, nullptr};
});
return res;
}
bool remove(K key) {
bool res = false;
update_impl(&root, key, [&](Node **at) {
if (!(*at))
return;
res = true;
release_from_tree(at);
});
return res;
}
V lookup(K key) {
Node* at = try_lookup(root, key);
assert(at);
return at->val;
}
static void printImpl(Node* n, int indent, std::vector<bool>& has_edge) {
if (!n)
return;
auto addIndent = [&]{
for (int i = 0; i < indent; i++) {
if (i < (indent - 1)) {
if (!has_edge[i])
std::cout << " ";
else
std::cout << "| ";
} else {
if (!has_edge[i])
std::cout << ",-";
else
std::cout << "`-";
has_edge[i] = !has_edge[i];
}
}
};
has_edge.resize(indent + 1, 0);
has_edge[indent] = 0;
printImpl(n->left, indent + 1, has_edge);
addIndent(); std::cout << n->depth << " [" << n->key << "] = " << n->val << std::endl;
has_edge[indent] = 1;
printImpl(n->right, indent + 1, has_edge);
}
void print() {
std::vector<bool> has_edge;
printImpl(root, 0, has_edge);
std::cout << std::endl;
}
int verify_impl(Node* n) {
if (!n)
return 0;
int l = verify_impl(n->left);
int r = verify_impl(n->right);
int d = std::max(l, r) + 1;
assert(n->depth == d);
return d;
}
void verify_depth() {
verify_impl(root);
}
};
int main() {
BalancedBinaryTree<int, int> tree;
int size = 20;
for (int i = 0; i < size; i++) {
tree.insert(i, i);
tree.print();
tree.verify_depth();
for (int k = 0; k <= i; k++)
assert(tree.lookup(k) == k);
}
for (int i = 0; i < size; i++) {
tree.remove(i);
tree.print();
tree.verify_depth();
for (int k = i + 1; k < size; k++)
assert(tree.lookup(k) == k);
}
}
It took me about 1h to write the code, and then I cleanup the code to make it presentable. Would that be acceptable dynamically balanced binary tree even if its its not a RB or AVL tree ?
RB or AVLaren't the most simple balanced binary search trees. \$\endgroup\$