/* Red-Black Tree implementation. Since the Red-Black Tree is quite an advanced data structure it's easy to make mistakes, this tree passed my tests but there might be some edge cases that I missed. Run on windows for colored output :) */ #include #include #include #include #include #define verify_tree #ifdef verify_tree #include #endif // verify_tree namespace dennis { #ifdef _WIN32 #include namespace System { void setConsoleColor(char c) { SetConsoleTextAttribute(GetStdHandle(STD_OUTPUT_HANDLE), c); } } #endif // _WIN32 template class RBTree { private: using size_type = std::size_t; struct Node; using NodePtr = std::shared_ptr; enum class Color : bool { Red, Black }; // Non-existing nodes are considered black, newly constructed nodes however are red. struct Node { Color color = Color::Red; NodePtr parent = nullptr; NodePtr left = nullptr; NodePtr right = nullptr; _Key key; Node() = default; Node(const _Key& key, NodePtr parent = nullptr): parent(parent), key(key){} }; public: // Insertion just like any other binary search tree, with the exception of a repair function at the end size_type size() const noexcept { return mSize; } bool empty() const noexcept { return mSize == 0; } bool insert(const _Key& key) { NodePtr* ptr = &mRoot; NodePtr* parent = nullptr; while(*ptr) { if(key < (*ptr)->key) { parent = ptr; ptr = &(*ptr)->left; } else if(key > (*ptr)->key) { parent = ptr; ptr = &(*ptr)->right; } else { return false; } } (*ptr).reset(new Node(key, parent ? *parent : nullptr)); ++mSize; repairTree(*ptr); return true; } bool remove(const _Key& key) { NodePtr n = find(key); if(n) remove(n); #ifdef verify_tree assert(assertTree()); #endif // verify_tree return n != nullptr; } // Yes, I know you can still call delete on the returned key and cause problems, // but that's a problem for another time :) std::pair search(const _Key& key) { NodePtr n = find(key); bool hasN = n != nullptr; return std::make_pair((hasN ? &n->key : nullptr),hasN); } #ifdef verify_tree int assertTree() { if(isRed(mRoot)) { std::cerr << "Root is black violation\n"; } return assertTree(mRoot); } #endif // verify_tree void inorderPrint() { inorderPrint(mRoot); } private: void inorderPrint(const NodePtr& n) { if(n) { inorderPrint(n->left); #ifdef _WIN32 System::setConsoleColor( isRed(n) ? 12 : 8 ); #endif // _WIN32 std::cout << n->key << ' '; #ifdef _WIN32 System::setConsoleColor( 7 ); #endif // _WIN32 inorderPrint(n->right); } } NodePtr find(const _Key& key) { NodePtr n = mRoot; while(n && key != n->key) { if(key < n->key) n = n->left; else if(key > n->key) n = n->right; } return n; } NodePtr findmin(const NodePtr& n) { NodePtr current = n; while(current->left) // Since the nodes on the left are smaller current = current->left; return current; } void replace(const NodePtr &n, const NodePtr &replacement = nullptr) { if(n->parent) // If n has a parent we have to update their connections as well { if(n == n->parent->left) // Child is on the left side, so update the parent's left link n->parent->left = replacement; else // Otherwise it must be on the right side n->parent->right = replacement; } else { // If n has no parent it means it is the root // So we have to update mRoot mRoot = replacement; } if(replacement) // If the replacement is not a nullptr we also have to link the replacement back to the parent replacement->parent = n->parent; } /* Remove has 3 cases to consider: 1) Deleting a node with no children: simply remove the node from the tree. 2) Deleting a node with 1 child: remove the node and replace it with its child. 3) Deleting a node with 2 children: call the node to be deleted D. Do not delete D. Instead, choose either its in-order predecessor node or its in-order successor node as replacement E. Copy the user values of E to D. If E does not have a child simply remove E from its previous parent G. If E has a child, say F, it is a right child. Replace E with F at E's parent. */ void remove(const NodePtr &n) { #ifdef verify_tree assert(n); #endif // verify_tree // 2 child case is the same as with the regular binary search tree removal if(n->left && n->right) // n has 2 children { NodePtr successor = findmin(n->right); // Find a successor = smallest key of its right child // Copy the key to n n->key = successor->key; // successor is guaranteed to have 1 or no children so it is a lot easier to delete that one remove(successor); } else { // Red Black Tree removal prerequisite is that it has only 1 or no child(ren). deleteNode(n); } } void deleteNode(const NodePtr &n) { // First find out which is the child NodePtr child = (n->right) ? n->right : n->left; // Child will be a nullptr if n has no children if(isBlack(n)) // If n is red we can safely remove it without any other issues { if(isRed(child)) // If the child of n is red simply color it black { child->color = Color::Black; } else // In other cases we have to do it the complicated way { /* n is a double black node. The tree can have 1 of 6 cases at this point, the task is to find the correct case that applies and resolve the problem for said case. The case to choose depends on the colors of other nodes. */ deleteCase1(n); } } replace(n, child); --mSize; #ifdef verify_tree assert(n.use_count() == 1); #endif // verify_tree } void deleteCase1(const NodePtr &n) { /* * Terminating case ( If this case applies we are done ) O = Double black (n), B = Black R = Red Before => After O => B / \ / \ */ // Case 1: Check if N is the root, if it is we are done. if(parent(n)) { deleteCase2(n); } } void deleteCase2(const NodePtr &n) { /* Non-Terminating ( After this case we move on to the next ) If this case applies the resulting tree violates property 5 ( all paths from any given node to its leaves have the same number of black nodes ) B => B / \ / \ O R R B / \ / \ / \ B B O B */ NodePtr s = sibling(n); #ifdef verify_tree assert(s); #endif // verify_tree if(isRed(s)) // Only need to check if sibling is red because in all other cases it's black { s->color = Color::Black; n->parent->color = Color::Red; if(n == n->parent->left) rotateLeft(n->parent); else rotateRight(n->parent); } deleteCase3(n); } void deleteCase3(const NodePtr &n) { /* Non-Terminating case If this case applies property 5 is still violated. B => O / \ / \ O B B R / \ / \ B B B B */ NodePtr s = sibling(n); #ifdef verify_tree assert(s); #endif // verify_tree if((isBlack(n->parent)) && (isBlack(s)) && (isBlack(s->left)) && (isBlack(s->right))) { s->color = Color::Red; deleteCase1(n->parent); } else { deleteCase4(n); } } void deleteCase4(const NodePtr &n) { /* * Terminating case, if this case applies we are done. R => B / \ / \ O B B R / \ / \ B B B B */ NodePtr s = sibling(n); #ifdef verify_tree assert(s); #endif // verify_tree if((isRed(n->parent)) && (isBlack(s)) && (isBlack(s->left)) && (isBlack(s->right))) { // Only have to swap colors of s and n's parent std::swap(s->color, n->parent->color); } else { deleteCase5(n); } } void deleteCase5(const NodePtr &n) { /* Non-terminating case B => B / \ / \ O B O B / \ \ R B R \ B */ NodePtr s = sibling(n); #ifdef verify_tree assert(s); #endif // verify_tree if((n == n->parent->left) && (isBlack(s->right)) && (isRed(s->left))) { std::swap(s->color, s->left->color); rotateRight(s); } else if((n == n->parent->right) && (isBlack(s->left)) && (isRed(s->right))) { std::swap(s->color, s->right->color); rotateLeft(s); } deleteCase6(n); } void deleteCase6(const NodePtr &n) { /* * Terminating case, last case to consider P = We don't care about the color P => P / \ / \ O B B B / \ / \ P R B P */ NodePtr s = sibling(n); #ifdef verify_tree assert(s); #endif // verify_tree std::swap(s->color, n->parent->color); if(n == n->parent->left) { s->right->color = Color::Black; rotateLeft(n->parent); } else { s->left->color = Color::Black; rotateRight(n->parent); } } NodePtr parent(const NodePtr &n) { return n ? n->parent : nullptr; } NodePtr grandParent(const NodePtr &n) { return parent(parent(n)); } NodePtr sibling(const NodePtr &n) { const NodePtr &p = parent(n); if(!p) return nullptr; return p->left == n ? p->right : p->left; } NodePtr uncle(const NodePtr &n) { return sibling(parent(n)); } void rotateLeft(NodePtr n) { /* Example: ( 1 to 4 may be (un)allocated ) (inside) p / \ n 4 / \ 1 x / \ 2 3 Result: (outside) p / \ x 4 / \ n 3 / \ 1 2 */ #ifdef verify_tree assert(n); #endif // verify_tree NodePtr x = n->right; // x will be 'disconnected' from the tree temporarily, so for it to not be lost make something point to it. n->right = x->left; // Make n->right point to node 2 // node 2 may not exist, so do a nullptr check before accessing it if(n->right != nullptr) n->right->parent = n; // now node 2's parent is n // make x's new parent n x->parent = n->parent; // Now we have to 'reconnect' x with the tree, // therefore x's parent has to point to x ( p is still pointing to n at this point ) // however it is possible that x is the root at this point so if that's the case there is no parent if(n->parent == nullptr) { mRoot = x; // x becomes the new emperor! } // Otherwise x can be on either the left or the right side of the parent, so find out which one else if(n == n->parent->left) // If n is on the left side of the parent { n->parent->left = x; } else // Otherwise it has to be on the right side { n->parent->right = x; } // Last step is to connect x and n together n->parent = x; x->left = n; // In left rotation, n will always end up on the left side } void rotateRight(NodePtr n) { /* Example: ( 1 to 4 may be (un)allocated ) (outside) n / \ x 4 / \ a 3 / \ 1 2 Result: (outside) x / \ a n / \ / \ 1 2 3 4 */ #ifdef verify_tree assert(n); #endif // verify_tree NodePtr x = n->left; // x will be 'disconnected' from the tree temporarily, so for it to not be lost make something point to it. n->left = x->right; // Make n->right point to node 2 // node 2 may not exist, so do a nullptr check before accessing it if(n->left != nullptr) n->left->parent = n; // now node 2's parent is n // make x's new parent n x->parent = n->parent; // Now we have to 'reconnect' x with the tree, // therefore x's parent has to point to x ( p is still pointing to n at this point ) // however it is possible that x is the root at this point so if that's the case there is no parent if(n->parent == nullptr) { mRoot = x; // x becomes the new emperor! } // Otherwise x can be on either the left or the right side of the parent, so find out which one else if(n == n->parent->left) // If n is on the left side of the parent { n->parent->left = x; } else // Otherwise it has to be on the right side { n->parent->right = x; } // Last step is to connect x and n together n->parent = x; x->right = n; // In right rotation, n will always end up on the right side } void repairTree(NodePtr n) { NodePtr p = parent(n); NodePtr g = grandParent(n); NodePtr u = uncle(n); // n is the root if(p == nullptr) // Property 2 (the root is black) { n->color = Color::Black; } else if(p->color == Color::Black) // Nothing to repair here { } else if(isRed(u)) // Property 5 ( all paths from any given node to its leaf nodes contain the same number of black nodes ) { // If u is red that means that u must exist ( as a non existing uncle is black ), // If u exist g must exist as well, // and if g exist p must exist. // Paint u and p black, and g red u->color = Color::Black; p->color = Color::Black; g->color = Color::Red; // g may now violate property 2 ( the root is black ) if it is the root // or property 4 ( Both children of every red node are black ) // so repair it again. repairTree(g); } else { // We have to rotate the parent node to its grandparent position. // But if n is on the 'inside' of the tree we have to move it to the 'outside' /* g <= This is called a left-right case (left from g and then right from p) / \ p u \ n <= n is on the inside becomes g / \ p u / n <= n is on the outside as well as the mirror case, which is called right-left case (right from g and then left from p) */ // p and g must exist here, however since u is black it may or may not exist if(g->left && n == g->left->right) // left-right case { rotateLeft(p); // Rotate the parent of n n = p; // p becomes the new n ( since the roles of n and p have swapped, n is now the parent of p ) } else if(g->right && n == g->right->left) // right-left case { rotateRight(p); n = p; } // Now that we fixed the n is on the 'inside' issue we can move the new parent to the grandparent position // But first we have to update the p and g node g = grandParent(n); p = parent(n); // Secondly find out in which direction to rotate in if(n == p->left) { rotateRight(g); } else { rotateLeft(g); } // Lastly swap the colors of p and g std::swap(p->color, g->color); } #ifdef verify_tree assert(assertTree()); #endif // verify_tree } bool isRed(const NodePtr& n) { return n && n->color == Color::Red; } bool isBlack(const NodePtr& n) { return !n || n->color == Color::Black; } #ifdef verify_tree int assertTree(const NodePtr &root) { if(!root) { return true; } NodePtr ln = root->left; NodePtr rn = root->right; if(root->color == Color::Red) { if((ln != nullptr && ln->color == Color::Red) || (rn != nullptr && rn->color == Color::Red)) { std::cerr << "Red violation\n"; return false; } } int lh = assertTree(ln); int rh = assertTree(rn); /* Invalid binary search tree */ if((ln != nullptr && ln->key >= root->key) || (rn != nullptr && rn->key <= root->key)) { std::cerr << "Binary tree violation\n"; return false; } /* Black height mismatch */ if(lh != 0 && rh != 0 && lh != rh) { std::cerr << "Black violation\n"; return false; } /* Only count black links */ if(lh != 0 && rh != 0) { return (root->color == Color::Red) ? lh : lh + 1; } else { return false; } } #endif // verify_tree private: NodePtr mRoot; size_type mSize = 0; }; } namespace Random { std::mt19937 mt(std::time(nullptr)); std::unordered_set randomize(dennis::RBTree& t, int min, int max, int count) { std::unordered_set s; while(count--) { int r = std::uniform_int_distribution (min, max)(mt); t.insert(r); s.insert(r); } return s; } } int main() { dennis::RBTree tree; std::unordered_set s = Random::randomize(tree, 0, 100, 15); std::cout << "In order print:\n"; tree.inorderPrint(); std::cout << '\n'; for(const auto& i : s) { std::cout << "\nRemoving " << i << '\n'; tree.remove(i); tree.inorderPrint(); std::cout << '\n'; } return 0; }