forked from vieter-v/libvieter
290 lines
9.0 KiB
C
290 lines
9.0 KiB
C
#include "vieter_tree_balancing.h"
|
|
#include <string.h>
|
|
|
|
bool vieter_tree_node_validate(vieter_tree_node *node,
|
|
uint64_t passed_black_nodes,
|
|
uint64_t expected_black_nodes) {
|
|
// Path to NULL children should contain the same amount of black nodes
|
|
if (node == NULL) {
|
|
return passed_black_nodes == expected_black_nodes;
|
|
}
|
|
|
|
if (vieter_tree_node_get(node, vieter_tree_node_black)) {
|
|
passed_black_nodes++;
|
|
}
|
|
|
|
// Either the node itself is black, or its children are both either NULL or
|
|
// black.
|
|
bool correctly_colored_children =
|
|
vieter_tree_node_get(node, vieter_tree_node_black) ||
|
|
((node->children[0] == NULL ||
|
|
vieter_tree_node_get(node->children[0], vieter_tree_node_black)) &&
|
|
(node->children[1] == NULL ||
|
|
vieter_tree_node_get(node->children[1], vieter_tree_node_black)));
|
|
bool right_child_flag_set =
|
|
node->parent == NULL ||
|
|
(vieter_tree_node_get(node, vieter_tree_node_right) ==
|
|
(node->parent->children[1] == node));
|
|
|
|
return correctly_colored_children && right_child_flag_set &&
|
|
vieter_tree_node_validate(node->children[0], passed_black_nodes,
|
|
expected_black_nodes) &&
|
|
vieter_tree_node_validate(node->children[1], passed_black_nodes,
|
|
expected_black_nodes);
|
|
}
|
|
|
|
/*
|
|
* This function should be rewritten to use tree rotations instead.
|
|
*/
|
|
vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
|
|
vieter_tree_node *parent = node->parent;
|
|
vieter_tree_node *grand_parent = parent->parent;
|
|
vieter_tree_node *root = grand_parent;
|
|
vieter_tree_node *children[2];
|
|
vieter_tree_node *grand_children[4];
|
|
|
|
uint64_t key_root = root->key;
|
|
void *data_root = root->data;
|
|
|
|
if (vieter_tree_node_get(node, vieter_tree_node_right)) {
|
|
children[0] = parent;
|
|
children[1] = node;
|
|
|
|
if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
|
|
root->key = parent->key;
|
|
root->data = parent->data;
|
|
|
|
parent->key = key_root;
|
|
parent->data = data_root;
|
|
|
|
grand_children[0] = grand_parent->children[0];
|
|
grand_children[1] = parent->children[0];
|
|
grand_children[2] = node->children[0];
|
|
grand_children[3] = node->children[1];
|
|
} else {
|
|
root->key = node->key;
|
|
root->data = node->data;
|
|
|
|
node->key = key_root;
|
|
node->data = data_root;
|
|
|
|
grand_children[0] = parent->children[0];
|
|
grand_children[1] = node->children[0];
|
|
grand_children[2] = node->children[1];
|
|
grand_children[3] = grand_parent->children[1];
|
|
}
|
|
} else {
|
|
children[0] = node;
|
|
children[1] = parent;
|
|
|
|
if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
|
|
root->key = node->key;
|
|
root->data = node->data;
|
|
|
|
node->key = key_root;
|
|
node->data = data_root;
|
|
|
|
grand_children[0] = grand_parent->children[0];
|
|
grand_children[1] = node->children[0];
|
|
grand_children[2] = node->children[1];
|
|
grand_children[3] = parent->children[1];
|
|
} else {
|
|
root->key = parent->key;
|
|
root->data = parent->data;
|
|
|
|
parent->key = key_root;
|
|
parent->data = data_root;
|
|
|
|
grand_children[0] = node->children[0];
|
|
grand_children[1] = node->children[1];
|
|
grand_children[2] = parent->children[1];
|
|
grand_children[3] = grand_parent->children[1];
|
|
}
|
|
}
|
|
|
|
vieter_tree_node_set(
|
|
root, vieter_tree_node_right,
|
|
vieter_tree_node_get(grand_parent, vieter_tree_node_right));
|
|
|
|
vieter_tree_node_set(children[0], vieter_tree_node_black, true);
|
|
vieter_tree_node_set(children[1], vieter_tree_node_black, true);
|
|
|
|
vieter_tree_node_set(root, vieter_tree_node_black, false);
|
|
|
|
vieter_tree_node_set_children(children[0], grand_children);
|
|
vieter_tree_node_set_children(children[1], grand_children + 2);
|
|
vieter_tree_node_set_children(root, children);
|
|
|
|
return root;
|
|
}
|
|
|
|
void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {
|
|
while (!vieter_tree_node_get(node, vieter_tree_node_black) &&
|
|
node->parent != NULL &&
|
|
!vieter_tree_node_get(node->parent, vieter_tree_node_black)) {
|
|
node = vieter_tree_node_balance(node);
|
|
}
|
|
|
|
// The root must always be black
|
|
if (node->parent == NULL) {
|
|
vieter_tree_node_set(node, vieter_tree_node_black, true);
|
|
}
|
|
}
|
|
|
|
vieter_tree_node *vieter_tree_node_rotate(vieter_tree_node *old_root,
|
|
bool dir) {
|
|
vieter_tree_node *new_root = old_root->children[1 - dir];
|
|
|
|
if (old_root->parent != NULL) {
|
|
vieter_tree_node_set_child(
|
|
old_root->parent, new_root,
|
|
vieter_tree_node_get(old_root, vieter_tree_node_right));
|
|
} else {
|
|
new_root->parent = NULL;
|
|
}
|
|
|
|
// Right rotation
|
|
if (dir) {
|
|
vieter_tree_node_set_child(old_root, new_root->children[1], false);
|
|
vieter_tree_node_set_child(new_root, old_root, true);
|
|
}
|
|
// Left rotation
|
|
else {
|
|
vieter_tree_node_set_child(old_root, new_root->children[0], true);
|
|
vieter_tree_node_set_child(new_root, old_root, false);
|
|
}
|
|
|
|
return new_root;
|
|
}
|
|
|
|
/*
|
|
* This function is currently implemented by very literally following the
|
|
* Wikipedia pseudocode. It's honestly not too bad, and I couldn't be bothered
|
|
* to properly research how to implement red-black removal (how is this so hard
|
|
* to find?)
|
|
*
|
|
* https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Removal_of_a_black_non-root_leaf
|
|
*/
|
|
vieter_tree_node *vieter_tree_node_remove_balanced(vieter_tree_node *node) {
|
|
vieter_tree_node *out;
|
|
|
|
if (node->parent == NULL) {
|
|
out = node->children[0] != NULL ? node->children[0] : node->children[1];
|
|
|
|
vieter_tree_node_free(node);
|
|
|
|
// This only happens when the root was the only element in the tree
|
|
if (out == NULL)
|
|
return out;
|
|
|
|
vieter_tree_node_set(out, vieter_tree_node_black, true);
|
|
out->parent = NULL;
|
|
|
|
return out;
|
|
}
|
|
|
|
// A red node can only have 0 or 2 children. The node we receive only has
|
|
// one child at most, so we know if it's red that it doesn't have any
|
|
// children. A black node that has a single (right) child can be replaced by
|
|
// this child, granted it becomes black as well. Either way, the node can be
|
|
// replaced by its right child (even if it's NULL).
|
|
if (!vieter_tree_node_get(node, vieter_tree_node_black) ||
|
|
node->children[1] != NULL) {
|
|
out = node->children[1];
|
|
|
|
vieter_tree_node_set_child(
|
|
node->parent, node->children[1],
|
|
vieter_tree_node_get(node, vieter_tree_node_right));
|
|
|
|
if (node->children[1] != NULL)
|
|
vieter_tree_node_set(node->children[1], vieter_tree_node_black, true);
|
|
|
|
vieter_tree_node_free(node);
|
|
|
|
return out;
|
|
}
|
|
|
|
// The complicated case is when we want to remove a black leaf
|
|
// https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Removal_of_a_black_non-root_leaf
|
|
|
|
out = node;
|
|
vieter_tree_node *parent = node->parent;
|
|
vieter_tree_node *sibling, *close_nephew, *distant_nephew;
|
|
bool dir = vieter_tree_node_get(node, vieter_tree_node_right);
|
|
parent->children[dir] = NULL;
|
|
|
|
vieter_tree_node_free(node);
|
|
|
|
goto start;
|
|
|
|
do {
|
|
dir = vieter_tree_node_get(node, vieter_tree_node_right);
|
|
|
|
start:
|
|
sibling = parent->children[1 - dir];
|
|
distant_nephew = sibling->children[1 - dir];
|
|
close_nephew = sibling->children[dir];
|
|
|
|
if (!vieter_tree_node_get(sibling, vieter_tree_node_black))
|
|
goto case3;
|
|
|
|
if (distant_nephew != NULL &&
|
|
!vieter_tree_node_get(distant_nephew, vieter_tree_node_black))
|
|
goto case6;
|
|
|
|
if (close_nephew != NULL &&
|
|
!vieter_tree_node_get(close_nephew, vieter_tree_node_black))
|
|
goto case5;
|
|
|
|
if (!vieter_tree_node_get(parent, vieter_tree_node_black))
|
|
goto case4;
|
|
|
|
// Case 2
|
|
vieter_tree_node_set(sibling, vieter_tree_node_black, false);
|
|
node = parent;
|
|
out = node;
|
|
} while ((parent = node->parent) != NULL);
|
|
|
|
// Case 1
|
|
return out;
|
|
|
|
case3:
|
|
out = vieter_tree_node_rotate(parent, dir);
|
|
vieter_tree_node_set(parent, vieter_tree_node_black, false);
|
|
vieter_tree_node_set(sibling, vieter_tree_node_black, true);
|
|
sibling = close_nephew;
|
|
distant_nephew = sibling->children[1 - dir];
|
|
|
|
if (distant_nephew != NULL &&
|
|
!vieter_tree_node_get(distant_nephew, vieter_tree_node_black))
|
|
goto case6;
|
|
|
|
close_nephew = sibling->children[dir];
|
|
|
|
if (close_nephew != NULL &&
|
|
!vieter_tree_node_get(close_nephew, vieter_tree_node_black))
|
|
goto case5;
|
|
|
|
case4:
|
|
vieter_tree_node_set(sibling, vieter_tree_node_black, false);
|
|
vieter_tree_node_set(parent, vieter_tree_node_black, true);
|
|
|
|
return out;
|
|
|
|
case5:
|
|
vieter_tree_node_rotate(sibling, 1 - dir);
|
|
vieter_tree_node_set(sibling, vieter_tree_node_black, false);
|
|
vieter_tree_node_set(close_nephew, vieter_tree_node_black, true);
|
|
distant_nephew = sibling;
|
|
sibling = close_nephew;
|
|
|
|
case6:
|
|
out = vieter_tree_node_rotate(parent, dir);
|
|
vieter_tree_node_set(sibling, vieter_tree_node_black,
|
|
vieter_tree_node_get(parent, vieter_tree_node_black));
|
|
vieter_tree_node_set(parent, vieter_tree_node_black, true);
|
|
vieter_tree_node_set(distant_nephew, vieter_tree_node_black, true);
|
|
|
|
return out;
|
|
}
|