Red Black Tree from OSD.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* parent;

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

    RBTreeNode* l() {
        return reinterpret_cast<RBTreeNode*>(this->left);
    }
    RBTreeNode* r() {
        return reinterpret_cast<RBTreeNode*>(this->right);
    }

    void pushBlack() {
        this->colour--;
        l()->colour++;
        r()->colour++;
    }
    void pullBlack() {
        this->colour++;
        l()->colour--;
        r()->colour--;
    }

    RBTreeNode* rotateLeft(RBTreeNode* u) {
        RBTreeNode* w = u->r();
        w->parent = u->parent;
        if (w->parent != nullptr) {
            if (w->parent->l() == u)
                w->parent->left = w;
            else
                w->parent->right = w;
        }
        u->right = w->left;
        if (u->right != nullptr)
            u->r()->parent = u;
        u->parent = w;
        w->left = u;
        if (u == this)
            return w;
        return this;
    }
    RBTreeNode* flipLeft(RBTreeNode* u) {
        char tmp = u->r()->colour;
        u->r()->colour = u->colour;
        u->colour = tmp;
        return rotateLeft(u);
    }

    RBTreeNode* rotateRight(RBTreeNode* u) {
        RBTreeNode* w = u->l();
        w->parent = u->parent;
        if (w->parent != nullptr) {
            if (w->parent->left == u)
                w->parent->left = w;
            else
                w->parent->right = w;
        }
        u->left = w->right;
        if (u->left != nullptr)
            u->l()->parent = u;
        u->parent = w;
        w->right = u;
        if (u == this)
            return w;
        return this;
    }
    RBTreeNode* flipRight(RBTreeNode* u) {
        char tmp = u->l()->colour;
        u->l()->colour = u->colour;
        u->colour = tmp;
        return rotateRight(u);
    }

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

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

        // Fixes red edges touching property and/or left-leaning property
        auto addFixup = [this](RBTreeNode* u)->RBTreeNode* {
            if (u == nullptr)
                throw runtime_error("addFixup: nullptr");
            RBTreeNode* root = this;
            while (u->colour == RED) {
                if (u == root) {
                    u->colour = BLACK;
                    break;
                }
                RBTreeNode* w = u->parent;
                if (w->left == nullptr || w->l()->colour == BLACK) { // ensure left-leaning
                    root = root->flipLeft(w);
                    u = w;
					w = u->parent;
                }
                if (w->colour == BLACK)
                    break; // no red-red edge = done
                RBTreeNode* g = w->parent; // grandparent of u
                //if (g == nullptr)
                //    continue;
                if (g->r()->colour == BLACK) {
                    root = root->flipRight(g);
                    break;
                }
                else {
                    g->pushBlack();
                    u = g;
                }
            }
            return root;
        };

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


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