From bee520a781e74a5c7bee1a55ca6f4539b85a232b Mon Sep 17 00:00:00 2001 From: Chewing_Bever Date: Wed, 1 Feb 2023 23:04:35 +0100 Subject: [PATCH] fix(tree): working balanced remove --- src/tree/vieter_tree.c | 2 +- src/tree/vieter_tree_balancing.c | 51 +++++++++++++++++++++++++++----- src/tree/vieter_tree_balancing.h | 31 +++++++++++++++++-- src/tree/vieter_tree_node.c | 17 +++++++---- src/tree/vieter_tree_node.h | 2 +- test/tree/test_balancing.c | 15 ++++++++++ test/tree/test_tree.c | 2 +- 7 files changed, 102 insertions(+), 18 deletions(-) diff --git a/src/tree/vieter_tree.c b/src/tree/vieter_tree.c index b066af1..8a00088 100644 --- a/src/tree/vieter_tree.c +++ b/src/tree/vieter_tree.c @@ -48,7 +48,7 @@ vieter_tree_error vieter_tree_remove(void **out, vieter_tree *tree, return vieter_tree_not_present; } - vieter_tree_error res = vieter_tree_node_remove(out, tree->root, key); + vieter_tree_error res = vieter_tree_node_remove(out, &tree->root, key); if (res != vieter_tree_ok) { return res; diff --git a/src/tree/vieter_tree_balancing.c b/src/tree/vieter_tree_balancing.c index 96b0c10..6e7dcd9 100644 --- a/src/tree/vieter_tree_balancing.c +++ b/src/tree/vieter_tree_balancing.c @@ -33,6 +33,9 @@ bool vieter_tree_node_validate(vieter_tree_node *node, 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; @@ -144,7 +147,9 @@ vieter_tree_node *vieter_tree_node_rotate(vieter_tree_node *old_root, if (dir) { vieter_tree_node_set_child(old_root, new_root->children[1], false); vieter_tree_node_set_child(new_root, old_root, true); - } else { + } + // 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); } @@ -152,7 +157,32 @@ vieter_tree_node *vieter_tree_node_rotate(vieter_tree_node *old_root, return new_root; } -void vieter_tree_node_remove_balanced(vieter_tree_node *node) { +/* + * 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 @@ -160,6 +190,8 @@ void vieter_tree_node_remove_balanced(vieter_tree_node *node) { // 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)); @@ -169,12 +201,13 @@ void vieter_tree_node_remove_balanced(vieter_tree_node *node) { vieter_tree_node_free(node); - return; + 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); @@ -209,13 +242,14 @@ void vieter_tree_node_remove_balanced(vieter_tree_node *node) { // Case 2 vieter_tree_node_set(sibling, vieter_tree_node_black, false); node = parent; + out = node; } while ((parent = node->parent) != NULL); // Case 1 - return; + return out; case3: - vieter_tree_node_rotate(parent, dir); + 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; @@ -234,7 +268,8 @@ case3: case4: vieter_tree_node_set(sibling, vieter_tree_node_black, false); vieter_tree_node_set(parent, vieter_tree_node_black, true); - return; + + return out; case5: vieter_tree_node_rotate(sibling, 1 - dir); @@ -244,9 +279,11 @@ case5: sibling = close_nephew; case6: - vieter_tree_node_rotate(parent, dir); + 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; } diff --git a/src/tree/vieter_tree_balancing.h b/src/tree/vieter_tree_balancing.h index 0d1acc5..e5dced0 100644 --- a/src/tree/vieter_tree_balancing.h +++ b/src/tree/vieter_tree_balancing.h @@ -3,6 +3,26 @@ #include "vieter_tree_node.h" +/* + * Check whether the given tree is a valid red-black tree. + * + * @param node root of the (sub)tree + * @param passed_black_nodes how many black nodes the recursion has already seen + * at this point. This should be initialized as 0 for the topmost call. + * @param expected_black_nodes the correct amount of black nodes to expect when + * a NULL child is encountered. + * @return true if the tree is valid, false otherwise + */ +bool vieter_tree_node_validate(vieter_tree_node *node, + uint64_t passed_black_nodes, + uint64_t expected_black_nodes); + +/* + * Balance a path of 3 nodes into a complete binary tree, with a red root and + * black children. + */ +vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node); + /* * Ensure the tree remains a valid red-black tree after having inserting the * node. @@ -12,10 +32,15 @@ void vieter_tree_node_balance_after_insert(vieter_tree_node *node); /* * Remove the given node, ensuring the tree remains a valid red-black tree. * - * @param node node to remove. This should have at most a single child, namely - * the right one. + * @param node node to remove. This should have at most a single child. If node + * isn't the root, this should be the right child, otherwise it can be either. + * @return root of the subtree that this function operated on. This can + * sometimes be the new root of the entire tree. */ -void vieter_tree_node_remove_balanced(vieter_tree_node *node); +vieter_tree_node *vieter_tree_node_remove_balanced(vieter_tree_node *node); +/* + * Perform a tree rotation of the subtree with the given root. + */ vieter_tree_node *vieter_tree_node_rotate(vieter_tree_node *old_root, bool dir); #endif diff --git a/src/tree/vieter_tree_node.c b/src/tree/vieter_tree_node.c index 3f09f3e..57aea05 100644 --- a/src/tree/vieter_tree_node.c +++ b/src/tree/vieter_tree_node.c @@ -109,16 +109,17 @@ vieter_tree_error vieter_tree_node_search(void **out, vieter_tree_node *root, return vieter_tree_ok; } -vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root, - uint64_t key) { +vieter_tree_error +vieter_tree_node_remove(void **out, vieter_tree_node **root_ptr, uint64_t key) { vieter_tree_node *target; - vieter_tree_error res = vieter_tree_node_search_node(&target, root, key); + vieter_tree_error res = vieter_tree_node_search_node(&target, *root_ptr, key); if (res != vieter_tree_ok) { return res; } *out = target->data; + vieter_tree_node *possible_new_root; if (target->children[0] != NULL && target->children[1] != NULL) { vieter_tree_node *replacement = target->children[1]; @@ -130,9 +131,15 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root, target->key = replacement->key; target->data = replacement->data; - vieter_tree_node_remove_balanced(replacement); + possible_new_root = vieter_tree_node_remove_balanced(replacement); } else { - vieter_tree_node_remove_balanced(target); + possible_new_root = vieter_tree_node_remove_balanced(target); + } + + if (possible_new_root == NULL) { + *root_ptr = NULL; + } else if (possible_new_root->parent == NULL) { + *root_ptr = possible_new_root; } return vieter_tree_ok; diff --git a/src/tree/vieter_tree_node.h b/src/tree/vieter_tree_node.h index f164704..fabee6b 100644 --- a/src/tree/vieter_tree_node.h +++ b/src/tree/vieter_tree_node.h @@ -33,7 +33,7 @@ vieter_tree_error vieter_tree_node_search_node(vieter_tree_node **out, vieter_tr vieter_tree_error vieter_tree_node_search(void **out, vieter_tree_node *root, uint64_t key); -vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root, uint64_t key); +vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node **root_ptr, uint64_t key); void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace, vieter_tree_node *replacement); diff --git a/test/tree/test_balancing.c b/test/tree/test_balancing.c index 5794ab9..cfda9d3 100644 --- a/test/tree/test_balancing.c +++ b/test/tree/test_balancing.c @@ -3,6 +3,9 @@ #include "vieter_tree_node.h" #include "vieter_tree_balancing.h" +// This uses the examples from wikipedia +// https://en.wikipedia.org/wiki/Tree_rotation + void test_rotate_right() { vieter_tree_node *a = vieter_tree_node_init(); vieter_tree_node *b = vieter_tree_node_init(); @@ -24,6 +27,12 @@ void test_rotate_right() { TEST_CHECK(new_root->children[1] == q); TEST_CHECK(new_root->children[1]->children[0] == b); TEST_CHECK(new_root->children[1]->children[1] == c); + + vieter_tree_node_free(a); + vieter_tree_node_free(b); + vieter_tree_node_free(p); + vieter_tree_node_free(c); + vieter_tree_node_free(q); } void test_rotate_left() { @@ -47,6 +56,12 @@ void test_rotate_left() { TEST_CHECK(new_root->children[1] == c); TEST_CHECK(new_root->children[0]->children[0] == a); TEST_CHECK(new_root->children[0]->children[1] == b); + + vieter_tree_node_free(a); + vieter_tree_node_free(b); + vieter_tree_node_free(p); + vieter_tree_node_free(c); + vieter_tree_node_free(q); } TEST_LIST = { diff --git a/test/tree/test_tree.c b/test/tree/test_tree.c index 6fdd430..2ab6870 100644 --- a/test/tree/test_tree.c +++ b/test/tree/test_tree.c @@ -42,7 +42,7 @@ void test_remove() { void *out; - for (uint64_t i = 0; i < 100; i++) { + for (uint64_t i = 0; i < 250; i++) { TEST_CHECK(vieter_tree_search(&out, tree, i) == vieter_tree_ok); TEST_CHECK(vieter_tree_remove(&out, tree, i) == vieter_tree_ok); TEST_CHECK(vieter_tree_validate(tree));