Red-Black tree

#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;
}