#include <iostream>
#include <vector>
using namespace std;
struct TNode {
int color; // 0 - black, 1 - red, 2 - double-black;
bool is_nil; // Is nil?
int val; // Node's value.
TNode *left; // Left sub-tree.
TNode *right; // Right sub-tree.
TNode *parent; // Parent of the node.
};
class RBtree {
private:
TNode *root;
TNode *new_node(int);
void set_node(TNode *, TNode *);
void check_nil(TNode *&);
int get_color(TNode *);
void set_color(TNode *&, int);
TNode *_find(TNode *, int);
TNode *_find_least(TNode *);
void rotate_left(TNode *);
void rotate_right(TNode *);
void _insert(int, TNode *&, TNode *);
void insert_fix(TNode *);
void _remove(TNode *&);
void remove_fix(TNode *);
public:
RBtree() {
root = nullptr;
}
TNode *find(int);
void insert(int);
void remove(int);
void print(TNode *root=nullptr, int offset=0);
};
// Private functions {
TNode *RBtree::new_node(int val) {
TNode *node = new TNode({1, 0, val, nullptr, nullptr, nullptr});
return node;
}
void RBtree::set_node(TNode *node, TNode *new_node) {
if (new_node != nullptr) {
new_node->parent = node->parent;
}
if (node->parent != nullptr) {
if (node == node->parent->left) {
node->parent->left = new_node;
} else {
node->parent->right = new_node;
}
} else {
this->root = new_node;
}
}
void RBtree::check_nil(TNode *&node) {
if (node->is_nil) {
if (node->parent != nullptr) {
if (node == node->parent->left) {
node->parent->left = nullptr;
} else {
node->parent->right = nullptr;
}
} else {
this->root = nullptr;
}
free(node);
}
}
int RBtree::get_color(TNode *node) {
if (node == nullptr) {
return 0;
}
return node->color;
}
void RBtree::set_color(TNode *&node, int color) {
if (node != nullptr) {
node->color = color;
}
}
TNode *RBtree::_find(TNode *root, int val) {
if (root == nullptr) {
return nullptr;
} else if (val == root->val) {
return root;
} else if (val < root->val) {
return _find(root->left, val);
} else {
return _find(root->right, val);
}
}
TNode *RBtree::_find_least(TNode *root) {
while (root->left != nullptr) {
root = root->left;
}
return root;
}
void RBtree::rotate_left(TNode *prev_root) {
TNode *new_root = prev_root->right;
prev_root->right = new_root->left;
if (new_root->left != nullptr) {
new_root->left->parent = prev_root;
}
new_root->parent = prev_root->parent;
if (prev_root->parent == nullptr) {
this->root = new_root;
} else if (prev_root == prev_root->parent->left) {
prev_root->parent->left = new_root;
} else {
prev_root->parent->right = new_root;
}
new_root->left = prev_root;
prev_root->parent = new_root;
}
void RBtree::rotate_right(TNode *prev_root) {
TNode *new_root = prev_root->left;
prev_root->left = new_root->right;
if (new_root->right != nullptr) {
new_root->right->parent = prev_root;
}
new_root->parent = prev_root->parent;
if (prev_root->parent == nullptr) {
this->root = new_root;
} else if (prev_root == prev_root->parent->left) {
prev_root->parent->left = new_root;
} else {
prev_root->parent->right = new_root;
}
new_root->right = prev_root;
prev_root->parent = new_root;
}
void RBtree::_insert(int val, TNode *&root, TNode *parent) {
if (root == nullptr) {
root = new_node(val);
root->parent = parent;
insert_fix(root);
return;
}
if (val < root->val) {
_insert(val, root->left, root);
} else if (val > root->val) {
_insert(val, root->right, root);
}
}
void RBtree::insert_fix(TNode *root) {
// Case 0. Root is always black.
// Case 1. Root's uncle is red.
// Case 2. Root's uncle is black, triangle.
// Case 3. Root's uncle is black, line.
while (root && get_color(root->parent) == 1) { // While parent is RED.
// Root's parent is on the left side.
if (root->parent == root->parent->parent->left) {
TNode *uncle = root->parent->parent->right;
if (get_color(uncle) == 1) { // Case 1. Uncle is RED.
set_color(root->parent, 0); // Case 1. Recolor parent.
set_color(uncle, 0); // Case 1. Recolor uncle.
set_color(root->parent->parent, 1); // Case 1. Recolor granp.
// Continue from grandparent, if exists.
root = root->parent ? root->parent->parent : nullptr;
} else {
if (root == root->parent->right) { // Case 2. Triangle.
root = root->parent; // Case 2.
rotate_left(root); // Case 2.
}
set_color(root->parent, 0); // Case 3. Line.
set_color(root->parent->parent, 1); // Case 3.
rotate_right(root->parent->parent); // Case 3. Rotate granp.
}
}
// Root's parent is on the right side.
else {
TNode *uncle = root->parent->parent->left;
if (get_color(uncle) == 1) { // Case 1. Uncle is RED.
set_color(root->parent, 0); // Case 1. Recolor parent.
set_color(uncle, 0); // Case 1. Recolor uncle.
set_color(root->parent->parent, 1); // Case 1. Recolor granp.
// Continue from grandparent, if exists.
root = root->parent ? root->parent->parent : nullptr;
} else {
if (root == root->parent->left) { // Case 2. Triangle.
root = root->parent; // Case 2.
rotate_right(root); // Case 2.
}
set_color(root->parent, 0); // Case 3. Line.
set_color(root->parent->parent, 1); // Case 3.
rotate_left(root->parent->parent); // Case 3. Rotate granp.
}
}
}
set_color(this->root, 0);
}
void RBtree::_remove(TNode *&root) {
if (root->left && root->right) { // Both children
TNode *successor = _find_least(root->right);
root->val = successor->val;
_remove(successor);
}
else if (!root->left && !root->right) { // Leaf
if (get_color(root) == 1) { // Root is red
set_node(root, nullptr);
free(root);
} else { // Root is black
root->is_nil = 1;
root->color = 2;
remove_fix(root);
}
}
else if (root->left || root->right) { // One child.
TNode *child;
if (root->left) { // Left child.
child = root->left;
} else { // Right child.
child = root->right;
}
set_node(root, child);
// Either root or it's child is red.
if (get_color(root) == 1 or get_color(child) == 1) {
set_color(child, 0);
free(root);
}
// Both root and it's child are black.
else {
set_color(child, 2);
remove_fix(child);
}
}
}
void RBtree::remove_fix(TNode *node) {
while (true) {
// Case 0. TERMINAL. Node is tree's root.
if (node->parent == nullptr) {
set_color(node, 0);
check_nil(node);
break;
}
// The node is not root, so we can find it's sibling.
TNode *sib = nullptr;
bool node_side = 0;
if (node == node->parent->left) {
sib = node->parent->right;
node_side = 0;
} else {
sib = node->parent->left;
node_side = 1;
}
// Case 1. Sibling is red.
if (get_color(sib) == 1) {
set_color(node->parent, 1);
set_color(sib, 0);
if (node_side == 0) {
rotate_left(node->parent);
} else {
rotate_right(node->parent);
}
continue;
}
// Case 2. Sibling is black, sib's children are black.
if (!get_color(sib) &&
!get_color(sib->left) && !get_color(sib->right)) {
set_color(sib, 1);
set_color(node, 0);
// Parent is red, terminal state.
if (get_color(node->parent) == 1) {
set_color(node->parent, 0);
check_nil(node);
break;
}
// Parent is black. Moving double-black to it.
else {
set_color(node->parent, 2);
check_nil(node);
node = node->parent;
}
continue;
}
// Case 3. Sibling is black, sib-inner:red, sib-outer: black.
// Case 4. Sibling is black, sib-inner:---, sib-outer: red.
if (get_color(sib) == 0) {
if (node_side == 0) { // Node is left child, sibling is right.
// Case 3.
if (get_color(sib->left) == 1 && get_color(sib->right) == 0) {
set_color(sib, 1);
set_color(sib->left, 0);
rotate_right(sib);
}
// Case 4. TERMINAL.
if (get_color(sib->right) == 1) {
set_color(sib->right, 0);
set_color(sib, get_color(sib->parent));
set_color(sib->parent, 0);
set_color(node, 0);
check_nil(node);
rotate_left(sib->parent);
break;
}
} else { // Node is right child, sibling is left.
// Case 3.
if (get_color(sib->right) == 1 && get_color(sib->left) == 0) {
set_color(sib, 1);
set_color(sib->right, 0);
rotate_left(sib);
}
// Case 4. TERMINAL.
if (get_color(sib->left) == 1) {
set_color(sib->left, 0);
set_color(sib, get_color(sib->parent));
set_color(sib->parent, 0);
set_color(node, 0);
check_nil(node);
rotate_right(sib->parent);
break;
}
}
}
}
set_color(this->root, 0);
}
// }
// Public functions {
TNode *RBtree::find(int val) {
return _find(this->root, val);
}
void RBtree::insert(int val) {
_insert(val, this->root, nullptr);
}
void RBtree::remove(int val) {
TNode *node = find(val);
if (node != nullptr) {
_remove(node);
}
}
void RBtree::print(TNode *root, int offset) {
if (root == nullptr && offset == 0) {
root = this->root;
}
if (root == nullptr) {
return;
}
this->print(root->right, offset + 1);
for (int i = 0; i < offset; i++) cout << " ";
if (root->color == 0) {
cout << "\033[1;34m";
} else if (root->color == 1) {
cout << "\033[1;31m";
} else if (root->color == 2) {
cout << "\033[1;30m";
}
if (root->is_nil) {
cout << "N";
} else {
cout << root->val;
}
cout << "\033[0m\n";
this->print(root->left, offset + 1);
}
// }
int main() {
RBtree tree;
vector<int> v = {8, 15, 23, 17, 2, 16, 7, 21, 22};
for (auto i : v) {
tree.insert(i);
}
tree.print();
int t, x;
while (1) {
cin >> t >> x;
if (t == 0) {
cout << "Inserting " << x << "...\n";
tree.insert(x);
} else {
cout << "Removing " << x << "...\n";
tree.remove(x);
}
cout << "----------------------------------------\n";
tree.print();
cout << "----------------------------------------\n\n\n";
}
return 0;
}
Comments
0 B
|👍
/👎