AVL Tree Implementation

AVL-tree.cpp

#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

#include <iostream>
#include <fstream>
#include <cstdio>
#include <vector>
#include <cmath>
#include <sstream>

enum Color
{
    Red,
    Green,
    Blue,
    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 Green:
                SetConsoleTextAttribute(hConsole, 10);
                break;
            case Blue:
                SetConsoleTextAttribute(hConsole, 9);
                break;
            default:
                SetConsoleTextAttribute(hConsole, 15);
                break;
            }
            return "";
        #endif
        #ifdef PLATFORM_LINUX
            // Linux color
            std::string color_code;
            switch (c)
            {
            case Red:
                color_code = "31";
                break;
            case Green:
                color_code = "32";
                break;
            case Blue:
                color_code = "34";
                break;
            default:
                color_code = "39";
                break;
            }
            return "\033[" + color_code + "m";
        #endif
    }
};

struct Node {
    int data;
    Node *left, *right;
    int height;
    int quantity = 0;
};

int height(Node *treeNode) {
    if(treeNode == NULL) return 0;
    else return treeNode->height;
}

int max(int a, int b) {
    return (a > b) ? a : b;
}

// Function that allocates a new node with the given key and NULL left and right pointers.
Node *newNode(int data) {
    Node *temp = new Node();
    temp->data = data;
    temp->left = temp->right = NULL;
    temp->height = 1;
    temp->quantity = 1;

    return temp;
}

// Function to perform right rotate
Node *rightRotate(Node *y){
    Node *x = y->left;
    Node *T2 = x->right;

    x->right = y;
    y->left = T2;

    y->height = max(height(y->left), height(y->right)) + 1;
    x->height = max(height(x->left), height(x->right)) + 1;

    return x;
}

// Function to perform right rotate
Node *leftRotate(Node *x){
    Node *y = x->right;
    Node *T2 = y->left;

    y->left = x;
    x->right = T2;

    x->height = max(height(x->left), height(x->right)) + 1;
    y->height = max(height(y->left), height(y->right)) + 1;

    return y;
}

// Function to get the tree balance factor
int getBalance(Node *treeNode) {
    if(treeNode == NULL) return 0;
    return height(treeNode->left) - height(treeNode->right);
}

// Recursive function to insert a key in the subtree rooted with node and returns the new root of the subtree.
Node *insert(Node *treeNode, int data){

    // Normal BST insertion
    if(treeNode == NULL) return newNode(data);

    if(data < treeNode->data)
    {
        treeNode->left = insert(treeNode->left, data); // If the key to be deleted is smaller than the root's key, then it lies in left subtree
    }
    else if(data > treeNode->data)
    {
        treeNode->right = insert(treeNode->right, data); // If the key to be deleted is greater than the root's key, then it lies in right subtree
    }
    else
    {
        treeNode->quantity += 1;
        return treeNode;
    }

    // Update height of this ancestor node
    treeNode->height = 1 + max(height(treeNode->left), height(treeNode->right)); // Get current Node Height

    // Get the balance factor of this ancestor node to check whether this node became unbalanced
    int balance = getBalance(treeNode); // Get the balance factor

    // If this node becomes unbalanced, then there are 4 cases
    if(balance > 1 && data < treeNode->left->data) return rightRotate(treeNode); // Left left case
    if(balance < -1 && data > treeNode->right->data) return leftRotate(treeNode); // Right right case
    if(balance > 1 && data > treeNode->left->data) { // Left right case
        treeNode->left = leftRotate(treeNode->left);
        return rightRotate(treeNode);
    }
    if(balance < -1 && data < treeNode->right->data) { // Right left case
        treeNode->right = rightRotate(treeNode->right);
        return leftRotate(treeNode);
    }

    return treeNode;
}

// return the node with minimum key value found in that tree. Note that the entire tree does not need to be searched.
Node *minValueNode(Node *node) {
    Node *current = node;
    while(current->left != NULL) current = current->left;
    return current;
}

Node *deleteNode(Node *treeNode, int data) {
    // Normal BST insertion
    if(treeNode == NULL) {
        return treeNode;
    }

    if(data < treeNode->data) treeNode->left = deleteNode(treeNode->left, data); // If the key to be deleted is smaller than the root's key, then it lies in left subtree
    else if(data > treeNode->data) treeNode->right = deleteNode(treeNode->right, data); // If the key to be deleted is greater than the root's key, then it lies in right subtree
    else {
        if((treeNode->left == NULL) || (treeNode->right == NULL)) {
            Node *temp = treeNode->left ? treeNode->left : treeNode->right;
            if(temp == NULL) { // No child case
                temp = treeNode;
                treeNode = NULL;
            }
            else // One child case
                *treeNode = *temp;
            free(temp);
        }
        else {
            // node with two children: Get the inorder
            // successor (smallest in the right subtree)
            Node *temp = minValueNode(treeNode->right);
            // Copy the inorder successor's
            // data to this node
            treeNode->data = temp->data;
            // Delete the inorder successor
            treeNode->right = deleteNode(treeNode->right, temp->data);
        }
    }

    if(treeNode == NULL) return treeNode;

    // Update height of this ancestor node
    treeNode->height = 1 + max(height(treeNode->left), height(treeNode->right));

    // Get the balance factor of this ancestor node to check whether this node became unbalanced
    int balance = getBalance(treeNode);

    if(balance > 1 && getBalance(treeNode->left) >= 0) return rightRotate(treeNode); // Left left rotate
    if(balance > 1 && getBalance(treeNode->left) < 0) { // Left right rotate
        treeNode->left = leftRotate(treeNode->left);
        return rightRotate(treeNode);
    }
    if(balance < -1 && getBalance(treeNode->left) >= 0) return leftRotate(treeNode); // Right right rotate
    if(balance < -1 && getBalance(treeNode->left) >= 0) { // Right left rotate
        treeNode->right = rightRotate(treeNode->right);
        return leftRotate(treeNode);
    }

    return treeNode;
}

// function to print preorder traversal of the tree. The function also prints height of every node
void preOrder(Node *root)
{
    if(root)
    {
        std::cout << root->data << " ";
        preOrder(root->left);
        preOrder(root->right);
    }
}

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

void BTreePrinter::print_node_internal(std::vector<Node*> nodes, int level, int max_level) {
    if (nodes.empty() || BTreePrinter::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::print_whitespaces(first_spaces);

    std::vector<Node*> 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 << 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;
            }

            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::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::print_whitespaces(first_spaces - i);
            if (nodes[j] == NULL)
            {
                BTreePrinter::print_whitespaces(edge_lines * 2 + i + 1);
                continue;
            }

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

            BTreePrinter::print_whitespaces(i * 2 - 1);

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

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

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

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

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

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

int main()
{
    int number, choice, choice1;
    std::string fileName;
    std::ifstream infile;
    int get_num;


    do {
        Node *root = NULL;
        std::cout << "******************************" << std::endl;
        std::cout << "1. Insert node manually" << std::endl;
        std::cout << "2. Insert node from file" << std::endl;
        std::cout << "0. Exit" << std::endl;
        std::cin >> choice;
        switch(choice) {
            case 1:
                do{
                    std::cout << "******************************" << std::endl;
                    std::cout << "OPTIONS" << std::endl;
                    std::cout << "1. Insert node" << std::endl;
                    std::cout << "2. Delete node" << std::endl;
                    std::cout << "3. Clear node" << std::endl;
                    std::cout << "4. Display node" << std::endl;
                    std::cout << "0. Back" << std::endl;
                    std::cin >> choice1;
                    switch(choice1) {
                        case 1:
                            std::cout << "Insert node: ";
                            std::cin >> number;
                            root = insert(root, number);
                            // printTree(root, 0, 0);
                            BTreePrinter::print_node(root);
                            break;

                        case 2:
                            std::cout << "Delete node: ";
                            std::cin >> number;
                            root = deleteNode(root, number);
                            BTreePrinter::print_node(root);
                            break;

                        case 3:
                            root = NULL;
                            std::cout << "Node cleared!" << std::endl;
                            break;

                        case 4:
                            std::cout << "Current Tree: " << std::endl;
                            BTreePrinter::print_node(root);
                            break;

                        case 0:
                            break;

                        default:
                            std::cout << "Wrong input" << std::endl;
                    }
                }
                    while(choice1 != 0);
                break;

            case 2:
                std::cout << "Please enter the file name: ";
                std::cin >> fileName;
                infile.open(fileName.c_str());
                if(!infile.is_open()) {
                    std::cout << "File failed to open" << std::endl;
                    return 0;
                }

                while(infile >> get_num) {
                    root = insert(root, get_num);
                }

                // printTree(root, 0, 0);
                BTreePrinter::print_node(root);

                break;
            case 0:
                break;

            default:
                std::cout << "Wrong input" << std::endl;
        }
    }
    while (choice != 0);

    return 0;
}