RBTree Solution.cpp

/*
Given an integer array nums where the elements are sorted in ascending order, convert it to a height-balanced binary search tree.

A height-balanced binary tree is a binary tree in which the depth of the two subtrees of every node never differs by more than one.
*/

#include <iostream>
#include <vector>
#include <queue>
#include <functional>

using namespace std;



struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};


/*
// Recursion method
class Solution {
public:
    static TreeNode* sortedArrayToBST(vector<int>& nums) {
        function<TreeNode*(int left, int right)> helper = [&](int left, int right)->TreeNode* {
            if (left > right)
                return nullptr;

            int mid = (right + left) / 2;
            if ((right + left) % 2 == 1)
                mid++;
            TreeNode* root = new TreeNode(nums[mid]);
            root->left = helper(left, mid-1);
            root->right = helper(mid+1, right);
            return root;
        };
        return helper(0, nums.size() - 1);
    }
};
*/




// RedBlack tree

char RED = 0;
char BLACK = 1;
struct RBTreeNode : TreeNode {
    char colour;

    RBTreeNode() : TreeNode(), colour(BLACK) {}
    RBTreeNode(int x) : TreeNode(x), colour(BLACK) { TreeNode(x);  }
    RBTreeNode(int x, TreeNode *left, TreeNode *right) : TreeNode(x, left, right), colour(BLACK) {}

    void pushBlack() {
        this->colour--;
        dynamic_cast<RBTreeNode*>(left)->colour++;
        this->right->colour++;
    }
    void pullBlack() {
        this->colour++;
        this->left->colour--;
        this->right->colour--;
    }

    static RBTreeNode* rotateLeft(RBTreeNode* u) {
        RBTreeNode* tmp = u->right;
        u->right = tmp->left;
        tmp->left = u;
        return tmp;
    }
    static RBTreeNode* flipLeft(RBTreeNode* u) {
        char tmp = u->right->colour;
        u->right->colour = u->colour;
        u->colour = tmp;
        return rotateLeft(u);
    }

    static RBTreeNode* rotateRight(RBTreeNode* u) {
        RBTreeNode* tmp = u->left;
        u->left = tmp->right;
        tmp->right = u;
        return tmp;
    }
    static RBTreeNode* flipRight(RBTreeNode* u) {
        char tmp = u->left->colour;
        u->left->colour = u->colour;
        u->colour = tmp;
        return rotateRight(u);
    }

    bool add(int x) {
        auto findLast = [](RBTreeNode* root, int x)->RBTreeNode* {
            RBTreeNode* iter = root, * prev = nullptr;
            while (iter != nullptr) {
                prev = root;
                int diff = x - iter->val;
                if (diff < 0)
                    iter = iter->left;
                else if (diff > 0)
                    iter = iter->right;
                else
                    return iter;
            }
            return prev;
        };

        auto addHelper = [this, &findLast](RBTreeNode* u)->bool {
            RBTreeNode* p = findLast(this, u->val);
            int diff = u->val - p->val;
            if (diff < 0)
                p->left = u;
            else if (diff > 0)
                p->right = u;
            else
                return false;
            return true;
        };

        // Fixes red edges touching property and/or left-leaning property
        auto addFixup = [this](RBTreeNode* u) {
            if (u == nullptr)
                throw runtime_error("addFixup: nullptr");

            // Create path to u so can get grandparent
            vector<RBTreeNode*> path;
            RBTreeNode* iter = this;
            while (iter != nullptr) {
                path.push_back(iter);
                int diff = u->val - iter->val;
                if (diff < 0)
                    iter = iter->left;
                else if (diff > 0)
                    iter = iter->right;
                else
                    break;
            }
            path.pop_back();
            if (iter != u)
                throw runtime_error("addFixup: iter is broken");

            while (u->colour == RED) {
                if (u == this) {
                    u->colour = BLACK;
                    return;
                }
                RBTreeNode* w = *(path.end() - 1); // parent of u
                if (w->left->colour == BLACK) { // ensure left-leaning
                    w = flipLeft(w);
                    u = w;
					w = *(path.end() - 1);
                }
                if (w->colour == BLACK)
                    return; // no red-red edge = done
                RBTreeNode* g = *(path.end() - 2); // grandparent of u
                if (g->right->colour == BLACK) {
                    flipRight(g);
                    return;
                }
                else {
                    g->pushBlack();
                    u = g;
                }
            }
        };

        RBTreeNode* u = new RBTreeNode(x);
        u->colour = RED;
        bool added = addHelper(u);
        if (added)
            addFixup(u);
        else
            delete u;
        return added;
    }
};


class Solution {
public:
    static TreeNode* sortedArrayToBST(vector<int>& nums) {
        if (nums.size() < 1)
            return nullptr;
        RBTreeNode* root = new RBTreeNode(nums[0]);
        for (int i = 1; i < nums.size(); i++)
            root->add(nums[i]);
        return root;
    }
};


void printTree(TreeNode* node) {
    if (node == nullptr)
        return;
    queue<TreeNode*> q, q2;
    q.push(node);
    while (!q.empty()) {
        TreeNode* u = q.front();
        q.pop();
        cout << u->val << " ";
        if (u->left != nullptr)
            q2.push(u->left);
        if (u->right != nullptr)
            q2.push(u->right);
        if (q.empty()) {
            q = q2;
            q2 = queue<TreeNode*>();
            cout << endl;
        }
    }
}


int main() {
    vector<int> v = { -10, -3, 0, 5, 9 };
    TreeNode* ret = Solution::sortedArrayToBST(v);
    printTree(ret);
    return 0;
}