Red-Black Tree Implementation

RBTree.cpp

#include <algorithm>
#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#include <string>
#include <sstream>

#if defined(_WIN32)
    #define PLATFORM_WINDOWS // Windows
#elif defined(_WIN64)
    #define PLATFORM_WINDOWS // Windows
#elif defined(__CYGWIN__) && !defined(_WIN32)
    #define PLATFORM_WINDOWS // Windows (Cygwin POSIX under Microsoft Window)
#elif defined(__linux__)
    #define PLATFORM_LINUX // Debian, Ubuntu, Gentoo, Fedora, openSUSE, RedHat, Centos and other
#endif

#ifdef PLATFORM_WINDOWS
    #include <windows.h>
#endif

enum Color {
    Red,
    Black,
    Default,
};

class Console {
    public:
    static std::string color(Color c)
    {
        #ifdef PLATFORM_WINDOWS
            // Windows color
            HANDLE  hConsole = GetStdHandle(STD_OUTPUT_HANDLE);

            switch (c)
            {
            case Red:
                SetConsoleTextAttribute(hConsole, 12);
                break;
            case Black:
                SetConsoleTextAttribute(hConsole, 10);
                break;
            default:
                SetConsoleTextAttribute(hConsole, 14);
                break;
            }
            return "";
        #endif
        #ifdef PLATFORM_LINUX
            // Linux color
            std::string color_code;
            switch (c)
            {
            case Red:
                color_code = "31";
                break;
            case Black:
                color_code = "32";
                break;
            default:
                color_code = "39";
                break;
            }
            return "\033[" + color_code + "m";
        #endif
    }
};

template <typename T>
class Node {
    public:
        T data;
        Color color;
        Node<T>* left;
        Node<T>* right;
        Node<T>* parent;
        int quantity;

        Node(T);
};

template <typename T>
Node<T>::Node(T data)
{
    this->data = data;
    this->quantity = 1;
    this->color = Red;
    this->left = NULL;
    this->parent = NULL;
    this->right = NULL;
}

template <typename T>
class BTreePrinter {
    public:
        static void print_node(Node<T>* root);
    
    private:
        static void print_node_internal(std::vector<Node<T>*> nodes, int level, int max_level);
        static void print_whitespaces(int count);
        static int max_level(Node<T> *node);
        static bool is_all_elements_null(std::vector<Node<T>*> list);
};

template <typename T>
void BTreePrinter<T>::print_node_internal(std::vector<Node<T>*> nodes, int level, int max_level) {
    if (nodes.empty() || BTreePrinter<T>::is_all_elements_null(nodes))
        return;

    int floor = max_level - level;
    int edge_lines = (int) std::pow(2, (std::max(floor - 1, 0)));
    int first_spaces = (int) std::pow(2, (floor)) - 1;
    int between_spaces = (int) std::pow(2, (floor + 1)) - 1;
    int additional_space = 0;

    BTreePrinter<T>::print_whitespaces(first_spaces);

    std::vector<Node<T>*> new_nodes;

    for (int i = 0; i < nodes.size(); i++) {
        additional_space = 0;

        if (nodes[i] != NULL) {
            // C++11
            // std::string printed_data = std::to_string(nodes[i]->data);
            
            // Non C++11
            std::stringstream stream;
            stream << nodes[i]->data;
            std::string printed_data = stream.str();
            
            std::cout << Console::color((nodes[i]->color == Black)?Black:Red) << printed_data;
            
            if (nodes[i]->quantity > 1)
            {
                // C++11
                // printed_data += std::to_string(nodes[i]->quantity);
                
                // Non C++11
                stream << nodes[i]->quantity;
                printed_data = stream.str();
                
                additional_space += 1;
                std::cout << ":" << nodes[i]->quantity;
            }

            std::cout << Console::color(Default);

            additional_space += printed_data.length() - 1;

            new_nodes.push_back(nodes[i]->left);
            new_nodes.push_back(nodes[i]->right);
        } else {
            new_nodes.push_back(NULL);
            new_nodes.push_back(NULL);
            std::cout << " ";
        }

        BTreePrinter<T>::print_whitespaces(between_spaces - additional_space);
    }

    std::cout << "\n";

    for (int i = 1; i <= edge_lines; i++)
    {
        for (int j = 0; j < nodes.size(); j++)
        {
            BTreePrinter<T>::print_whitespaces(first_spaces - i);
            if (nodes[j] == NULL)
            {
                BTreePrinter<T>::print_whitespaces(edge_lines * 2 + i + 1);
                continue;
            }

            if (nodes[j]->left != NULL)
            {
                std::cout << "/";
            }
            else
            {
                BTreePrinter<T>::print_whitespaces(1);
            }

            BTreePrinter<T>::print_whitespaces(i * 2 - 1);

            if (nodes[j]->right != NULL)
            {
                std::cout << "\\";
            }
            else
            {
                BTreePrinter<T>::print_whitespaces(1);
            }
            
            BTreePrinter<T>::print_whitespaces(edge_lines * 2 - i);
        }
        std::cout << "\n";
    }

    BTreePrinter<T>::print_node_internal(new_nodes, level + 1, max_level);
}

template <typename T>
void BTreePrinter<T>::print_whitespaces(int count) {
    for (int i = 0; i < count; i++)
        std::cout << " ";
}

template <typename T>
void BTreePrinter<T>::print_node(Node<T>* root) {
    std::vector<Node<T>*> nodes;
    nodes.push_back(root);
    BTreePrinter<T>::print_node_internal(nodes, 1, max_level(root));
}

template <typename T>
bool BTreePrinter<T>::is_all_elements_null(std::vector<Node<T>*> list) {
    for(int i = 0; i < list.size(); i++) {
        if (list[i] != NULL)
            return false;
    }
    return true;
}

template <typename T>
int BTreePrinter<T>::max_level(Node<T> *node) {
    if (node == NULL)
        return 0;
    return std::max(BTreePrinter<T>::max_level(node->left), BTreePrinter<T>::max_level(node->right)) + 1;
}

template <typename T>
class RBTree {
    public:
        RBTree();
        Node<T>* root_node;

        void insert(T);
        void print();
        void remove(Node<T>*);

    private:
        void insert_fix(Node<T>*);
        void delete_fix(Node<T>*);
        void transplant(Node<T>*, Node<T>*);
        Node<T>* minimum_node(Node<T>*);
        void left_rotate(Node<T>*);
        void right_rotate(Node<T>*);
};

template <typename T>
RBTree<T>::RBTree()
{
    this->root_node = NULL;
}

// Page 313
template <typename T>
void RBTree<T>::left_rotate(Node<T>* node)
{
    Node<T>* right_node = node->right;
    node->right = right_node->left;

    if (right_node->left != NULL)
    {
        right_node->left->parent = node;
    }

    right_node->parent = node->parent;

    if (node->parent == NULL)
    {
        this->root_node = right_node;
    }
    else if (node == node->parent->left)
    {
        node->parent->left = right_node;
    }
    else
    {
        node->parent->right = right_node;
    }
    
    right_node->left = node;
    node->parent = right_node;
}

// Symetric with left_rotate
template <typename T>
void RBTree<T>::right_rotate(Node<T>* node)
{
    Node<T>* left_node = node->left;
    node->left = left_node->right;

    if (left_node->right != NULL)
    {
        left_node->right->parent = node;
    }

    left_node->parent = node->parent;

    if (node->parent == NULL)
    {
        this->root_node = left_node;
    }
    else if (node == node->parent->right)
    {
        node->parent->right = left_node;
    }
    else
    {
        node->parent->left = left_node;
    }
    
    left_node->right = node;
    node->parent = left_node;
}

// Page 315
template <typename T>
void RBTree<T>::insert(T data)
{
    Node<T>* parent_current_node = NULL;
    Node<T>* current_node = this->root_node;
    Node<T>* inserted_node = new Node<T>(data);

    while (current_node != NULL)
    {
        parent_current_node = current_node;
        if (inserted_node->data < current_node->data)
        {
            current_node = current_node->left;
        }
        else if (inserted_node->data > current_node->data)
        {
            current_node = current_node->right;
        }
        else
        {
            current_node->quantity += 1;
            return;
        }   
    }

    inserted_node->parent = parent_current_node;

    if (parent_current_node == NULL)
    {
        this->root_node = inserted_node;
    }
    else if (inserted_node->data < parent_current_node->data)
    {
        parent_current_node->left = inserted_node;
    }
    else if (inserted_node->data > parent_current_node->data)
    {
        parent_current_node->right = inserted_node;
    }

    this->insert_fix(inserted_node);
}

// Page 316
template <typename T>
void RBTree<T>::insert_fix(Node<T>* inserted_node)
{
    while (inserted_node != this->root_node && inserted_node->parent->color == Red)
    {
        if (inserted_node->parent == inserted_node->parent->parent->left)
        {
            Node<T>* temp_node = inserted_node->parent->parent->right;
            if (temp_node->color == Red)
            {
                inserted_node->parent->color = Black;
                temp_node->color = Black;
                inserted_node->parent->parent->color = Red;
                inserted_node = inserted_node->parent->parent;
            }
            else
            {
                if (inserted_node == inserted_node->parent->right)
                {
                    inserted_node = inserted_node->parent;
                    this->left_rotate(inserted_node);
                }
                inserted_node->parent->color = Black;
                inserted_node->parent->parent->color = Red;
                this->right_rotate(inserted_node->parent->parent);
            }
        }
        // Symetric with the if clause
        else
        {
            Node<T>* temp_node = inserted_node->parent->parent->left;
            if (temp_node->color == Red)
            {
                inserted_node->parent->color = Black;
                temp_node->color = Black;
                inserted_node->parent->parent->color = Red;
                inserted_node = inserted_node->parent->parent;
            }
            else
            {
                if (inserted_node == inserted_node->parent->left)
                {
                    inserted_node = inserted_node->parent;
                    this->left_rotate(inserted_node);
                }
                inserted_node->parent->color = Black;
                inserted_node->parent->parent->color = Red;
                this->right_rotate(inserted_node->parent->parent);
            }
        }
        
    }

    this->root_node->color = Black;
}

template <typename T>
void RBTree<T>::print()
{
    BTreePrinter<T>::print_node(this->root_node);
}

template <typename T>
Node<T>* RBTree<T>::minimum_node(Node<T>* node)
{
    Node<T>* current_node = node;
    
    while (current_node->left != NULL)
    {
        current_node = current_node->left;
    }

    return current_node;
}

template <typename T>
void RBTree<T>::transplant(Node<T>* node_1, Node<T>* node_2)
{
    if (node_1->parent == NULL)
    {
        this->root_node = node_2;
    }
    else if (node_1 == node_1->parent->left)
    {
        node_1->parent->left = node_2;
    }
    else
    {
        node_1->parent->right = node_2;    
    }

    node_2->parent = node_1->parent;
}

template <typename T>
void RBTree<T>::remove(Node<T>* node)
{
    if (node->quantity > 1)
    {
        node->quantity -= 1;
        return;
    }

    Node<T>* node_temp = node;
    Color node_temp_original_color = node->color;
    Node<T>* node_track = NULL;

    if (node->left == NULL)
    {
        node_track = node->right;
        this->transplant(node, node->right);
    }
    else if (node->right == NULL)
    {
        node_track = node->left;
        this->transplant(node, node->left);
    }
    else
    {
        node_temp = this->minimum_node(node->right);
        node_temp_original_color = node_temp->color;
        node_track = node_temp->right;

        if (node_temp->parent == node)
        {
            node_track->parent = node;
        }
        else
        {
            this->transplant(node_temp, node_temp->right);
            node_temp->right = node->right;
            node_temp->right->parent = node_temp;
        }
        this->transplant(node, node_temp);

        node_temp->left = node->left;
        node_temp->left->parent = node_temp;
        node_temp->color = node->color;
    }

    if (node_temp_original_color == Black)
        this->delete_fix(node_track);
}

// Page 326
template <typename T>
void RBTree<T>::delete_fix(Node<T>* node)
{
    while (node != this->root_node && node->color == Black)
    {
        if (node == node->parent->left)
        {
            Node<T>* node_temp = node->parent->right;
            if (node_temp->color == Red)
            {
                node_temp->color == Black;
                node->parent->color = Red;
            }

            if (node_temp->left->color == Black && node_temp->right->color == Black)
            {
                node_temp->color = Red;
                node = node->parent;
            }
            else if (node_temp->right->color == Black)
            {
                node_temp->left->color = Black;
                node_temp->color = Red;
                this->right_rotate(node_temp);
                node_temp = node->parent->right;
            }
            node_temp->color = node->parent->color;
            node->parent->color = Black;
            node_temp->right->color = Black;
            this->left_rotate(node->parent);
            node = this->root_node;
        }
        else
        {
            Node<T>* node_temp = node->parent->left;
            if (node_temp->color == Red)
            {
                node_temp->color == Black;
                node->parent->color = Red;
            }

            if (node_temp->right->color == Black && node_temp->right->color == Black)
            {
                node_temp->color = Red;
                node = node->parent;
            }
            else if (node_temp->left->color == Black)
            {
                node_temp->right->color = Black;
                node_temp->color = Red;
                this->right_rotate(node_temp);
                node_temp = node->parent->left;
            }
            node_temp->color = node->parent->color;
            node->parent->color = Black;
            node_temp->left->color = Black;
            this->left_rotate(node->parent);
            node = this->root_node;
        }
        
    }
}

int main()
{
    RBTree<int> rbtree;
    rbtree.insert(11);
    rbtree.insert(2);
    rbtree.insert(14);
    rbtree.insert(15);
    rbtree.insert(1);
    rbtree.insert(7);
    rbtree.insert(5);
    rbtree.insert(8);
    rbtree.insert(4);
    rbtree.insert(8);
    rbtree.print();
    rbtree.remove(rbtree.root_node->right);
    rbtree.print();
    return 0;
}