YangZX's implementation of red black tree

gistfile1.cpp

/* yangzx's red-black tree! */
#include <cstdio>
#include <algorithm>
#include <cstdlib>

using namespace std;

struct RBTree{
    struct RBNode{
        int val;
        RBNode *ch[2], *p;
        char color;
    };
    const static char red = -1;
    const static char black = 0;
    enum CHILD {L, R};

    RBNode *rt;
    RBNode *nil;
    RBTree() {
        nil = new RBNode();
        nil -> color = black;
        rt = nil;
    }
    ~RBTree(){
        deleteTree(rt);
        delete nil;
    }

    int Dir(int a, int b)
    {
        if(a < b)
            return L;
        else
            return R;
    }

    void rotateLeftSon(RBNode* z, int l, int r)
    {
        RBNode *p = z -> p;
        RBNode *ls = z -> ch[l];
        RBNode *lrs = ls -> ch[r];
        ls -> ch[r] = z;
        z -> p = ls;
        z -> ch[l] = lrs;
        lrs -> p = z;
        ls -> p = p;
        if(p == nil){
            rt = ls;
        }else if(p -> ch[l] == z){
            p -> ch[l] = ls;
        }else{
            p -> ch[r] = ls;
        }
    }

    void fixup(RBNode* z)
    {
        // Loop invariants:
        // 1. z is red
        // 2. if z->p is root, then z->p is black
        // 3. if the tree violates any properties of the red black tree, then
        //    it violates at most one of them. And the violation is either "root 
        //    is black"(z is root and z is red in this case)  or "children of red
        //    node are black" (z and z->p is red in this case)
        while(z -> p -> color == red){
            int l = L, r = R;
            RBNode *p, *gp;
            p = z -> p;
            gp = p -> p;
            if(p == gp -> ch[R])
                swap(l, r);
            RBNode* y = gp -> ch[r];
            if(y -> color == red){
                y -> color = black;
                p -> color = black;
                gp -> color = red;
                z = gp;
            }else{
                if(z == p -> ch[r]){
                    z = p;
                    rotateLeftSon(z, r, l);
                }
                z -> p -> color = black;
                gp -> color = red;
                rotateLeftSon(gp, l, r);
            }
        }
        if(rt -> color == red)
            rt -> color = black;
    }

    void insert(int z)
    {
        RBNode *y = nil;
        RBNode *x = rt;
        while(x != nil){
            y = x;
            x = y -> ch[Dir(z, y -> val)];
        }
        RBNode* nz = new RBNode;
        nz -> val = z;
        nz -> ch[L] = nz -> ch[R] = nil;
        nz -> p = y;
        nz -> color = red;
        y -> ch[Dir(z, y -> val)] = nz;
        if(rt == nil)
            rt = nz;
        fixup(nz);
    }

    void foreach(void (*f)(int))
    {
        foreachNode(rt, f);
    }

    
    private:
    void deleteTree(RBNode* rt)
    {
        if(rt == nil)
            return;
        deleteTree(rt -> ch[L]);
        deleteTree(rt -> ch[R]);
        delete rt;
    }

    void foreachNode(RBNode* rt, void (*f)(int))
    {
        if(rt == nil)
            return;
        foreachNode(rt -> ch[L], f);
        f(rt -> val);
        foreachNode(rt -> ch[R], f);
    }
};


void print(int k)
{
    printf("%d ", k);
}

int main()
{
    RBTree t;
    int n;
    scanf("%d", &n);
    for(int i = 1; i <= n; i++){
        int k;
        scanf("%d", &k);
        t.insert(k);
    }
    t.foreach(print);
    printf("\n");
    return 0;
}