diff --git a/src/tree/balancing.c b/src/tree/balancing.c index 85ae1a3..b7466ea 100644 --- a/src/tree/balancing.c +++ b/src/tree/balancing.c @@ -1,3 +1,95 @@ #include "balancing.h" -void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {} +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 *grand_grand_parent = grand_parent->parent; + vieter_tree_node *root, *left, *right; + vieter_tree_node *children[4]; + + if (node->key < parent->key && parent->key < grand_parent->key) { + root = parent; + left = node; + right = grand_parent; + + children[0] = node->left; + children[1] = node->right; + children[2] = parent->right; + children[3] = grand_parent->right; + } else if (node->key < parent->key && parent->key > grand_parent->key) { + root = node; + left = grand_parent; + right = parent; + + children[0] = grand_parent->left; + children[1] = node->left; + children[2] = node->right; + children[3] = parent->right; + } else if (node->key > parent->key && parent->key < grand_parent->key) { + root = node; + left = parent; + right = grand_parent; + + children[0] = parent->left; + children[1] = node->left; + children[2] = node->right; + children[3] = grand_parent->right; + } else { + root = parent; + left = grand_parent; + right = node; + + children[0] = grand_parent->left; + children[1] = parent->left; + children[2] = node->left; + children[3] = node->right; + } + + vieter_tree_node_set(left, vieter_node_black); + vieter_tree_node_set(right, vieter_node_black); + vieter_tree_node_unset(root, vieter_node_black); + + left->left = children[0]; + left->right = children[1]; + + if (children[0] != NULL) + children[0]->parent = left; + + if (children[1] != NULL) + children[1]->parent = left; + + right->left = children[2]; + right->right = children[3]; + + if (children[2] != NULL) + children[2]->parent = right; + + if (children[3] != NULL) + children[3]->parent = right; + + root->left = left; + root->right = right; + left->parent = root; + right->parent = root; + + if (grand_grand_parent != NULL) { + vieter_tree_node_add_child(grand_grand_parent, root->key, root); + } else { + root->parent = NULL; + } + + return root; +} + +void vieter_tree_node_balance_after_insert(vieter_tree_node *node) { + while (!vieter_tree_node_get(node, vieter_node_black) && + node->parent != NULL && + !vieter_tree_node_get(node->parent, vieter_node_black)) { + node = vieter_tree_node_balance(node); + } + + // The root must always be black + if (node->parent == NULL) { + vieter_tree_node_set(node, vieter_node_black); + } +} diff --git a/src/tree/node.c b/src/tree/node.c index b36feed..d6c2b6d 100644 --- a/src/tree/node.c +++ b/src/tree/node.c @@ -1,4 +1,5 @@ #include "node.h" +#include "balancing.h" vieter_tree_node *vieter_tree_node_init() { return calloc(1, sizeof(vieter_tree_node)); @@ -55,6 +56,7 @@ vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key, new_node->data = data; vieter_tree_node_add_child(parent, key, new_node); + vieter_tree_node_balance_after_insert(new_node); return vieter_tree_ok; } @@ -95,20 +97,6 @@ vieter_tree_error vieter_tree_node_search(void **out, vieter_tree_node *root, return vieter_tree_ok; } -vieter_tree_node *vieter_tree_node_successor(vieter_tree_node *node) { - if ((node->left != NULL) ^ (node->right != NULL)) { - return node->left != NULL ? node->left : node->right; - } - - node = node->right; - - while (node->left != NULL) { - node = node->left; - } - - return node; -} - vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root, uint64_t key) { vieter_tree_node *target; diff --git a/src/tree/node.h b/src/tree/node.h index 2db2145..522ee59 100644 --- a/src/tree/node.h +++ b/src/tree/node.h @@ -37,4 +37,7 @@ void vieter_tree_node_set(vieter_tree_node *node, vieter_node_flag flag); void vieter_tree_node_unset(vieter_tree_node *node, vieter_node_flag flag); bool vieter_tree_node_get(vieter_tree_node *node, vieter_node_flag flag); +void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key, + vieter_tree_node *child); + #endif diff --git a/src/tree/tree.c b/src/tree/tree.c index c2c7aa7..d2579e6 100644 --- a/src/tree/tree.c +++ b/src/tree/tree.c @@ -16,6 +16,7 @@ vieter_tree_error vieter_tree_insert(vieter_tree *tree, uint64_t key, tree->root = vieter_tree_node_init(); tree->root->key = key; tree->root->data = data; + vieter_tree_node_set(tree->root, vieter_node_black); tree->size = 1; return vieter_tree_ok;