红黑树

/*
 * 红黑树
 *      1 提供基本的操作集：
 *          得到最大结点
 *          得到最小结点
 *          得到后继结点
 *          得到前驱结点
 *          插入结点
 *          删除结点
 *      2 可以扩展的操作
 *          输出一个有序列
 * 本例中没有对重复的元素进行控制
 *
 * NIL 这个思想是从《算法导论》中来的
 * 有了 NIL 所有的结点都是内结点，便于算法的操作
 *
 * Successor 算法需要明白两种情况，没有右孩子，有右孩子
 * InsertNode 算法只插入叶子
 * DeleteNode 算法不删除既有右孩子又有左孩子的结点，这种情况需要转化
 *
 */

#include <cstdio>

#define RED    1
#define BLACK  0
struct RBNode {
    int key;
    int color;
    RBNode *pParent;
    RBNode *pLChild;
    RBNode *pRChild;
};

class CRBTree {
public:
    CRBTree() {
        NIL = new RBNode;
        NIL->color = BLACK;
        NIL->pParent = NIL;
        NIL->pLChild = NIL;
        NIL->pRChild = NIL;

        root = NIL;
    }

    ~CRBTree() {
        delete NIL;
    }

    int GetMaxKey() {
        RBNode *node = GetMaxNode(root);
        if (node == NIL) return 0;
        else return node->key;
    }

    int GetMinKey() {
        RBNode *node = GetMinNode(root);
        if (node == NIL) return 0;
        else return node->key;
    }

    void InsertKey(int key) {
        RBNode *node = new RBNode;
        node->key = key;
        node->color = RED;
        node->pParent = NIL;
        node->pLChild = NIL;
        node->pRChild = NIL;

        InsertNode(node);
    }

    void DeleteKey(int key) {
        RBNode *node;
        while ((node=FindInTree(key)) != NIL) {
            node = DeleteNode(node);
            if (node != NIL) delete node;
        }
    }

    bool KeyInTree(int key) {
        return FindInTree(key) != NIL;
    }

private:
    RBNode *GetMaxNode(RBNode *node) {
        RBNode *p = node;
        while (p->pRChild != NIL) {
            p = p->pRChild;
        }
        return p;
    }

    RBNode *GetMinNode(RBNode *node) {
        RBNode *p = node;
        while (p->pLChild != NIL) {
            p = p->pLChild;
        }
        return p;
    }

    RBNode *FindInTree(int key) {
        RBNode *p = root;
        while (p!=NIL && p->key!=key) {
            if (p->key < key) p = p->pRChild;
            else p = p->pLChild;
        }
        return  p;
    }

    RBNode *Successor(RBNode *node) {
        if (node->pRChild != NIL) return GetMinNode(node->pRChild);

        RBNode *p = node;
        RBNode *pParent = p->pParent;
        while (pParent!=NIL && pParent->pRChild==p) {
            p = pParent;
            pParent = p->pParent;
        }

        return pParent;
    }

    RBNode *Predecessor(RBNode *node) {
        if (node->pLChild != NIL) return GetMaxNode(node->pLChild);

        RBNode *p = node;
        RBNode *pParent = p->pParent;
        while (pParent!=NIL && pParent->pLChild==p) {
            p = pParent;
            pParent = p->pParent;
        }
        return pParent;
    }

    void LeftRotate(RBNode *node) {
        // x y 都不是 NIL
        RBNode *x = node;
        RBNode *y = node->pRChild;

        y->pParent = x->pParent;
        if (x->pParent == NIL) root = y;
        else if (x->pParent->pLChild == x) x->pParent->pLChild = y;
        else x->pParent->pRChild = y;

        x->pRChild = y->pLChild;
        if (x->pRChild != NIL) x->pRChild->pParent = x;

        y->pLChild = x;
        x->pParent = y;
    }

    void RightRotate(RBNode *node) {
        // x y 都不是 NIL
        RBNode *x = node;
        RBNode *y = node->pLChild;
        y->pParent = x->pParent;
        if (x->pParent == NIL) root = y;
        else if (x->pParent->pLChild == x) x->pParent->pLChild = y;
        else x->pParent->pRChild = y;

        x->pLChild = y->pRChild;
        if (x->pLChild != NIL) x->pLChild->pParent = x;

        y->pRChild = x;
        x->pParent = y;
    }

    void InsertNode(RBNode *node) {
        RBNode *p = root;
        RBNode *pParent = NIL;

        while (p != NIL) {
            pParent = p;
            if (p->key < node->key) p = p->pRChild;
            else p = p->pLChild;
        }

        node->pParent = pParent;
        if (pParent == NIL) root = node;
        else if (pParent->key < node->key) pParent->pRChild = node;
        else pParent->pLChild = node;

        node->pLChild = NIL;
        node->pRChild = NIL;
        node->color = RED;

        InsertFixup(node);
    }

    void InsertFixup(RBNode *node) {
        RBNode *z = node;
        while (z->pParent->color==RED) {
            if (z->pParent==z->pParent->pParent->pLChild) {
                if (z->pParent->pParent->pRChild->color==RED) {
                    z->pParent->color =  BLACK;
                    z->pParent->pParent->pRChild->color= BLACK;
                    z->pParent->pParent->color = RED;
                    z = z->pParent->pParent;
                }
                else {
                    if (z->pParent->pRChild == z) {
                        z = z->pParent;
                        LeftRotate(z);
                    }

                    z->pParent->color = BLACK;
                    z->pParent->pParent->color = RED;
                    RightRotate(z->pParent->pParent);
                }
            }
            else {
                if (z->pParent->pParent->pLChild->color == RED) {
                    z->pParent->color = BLACK;
                    z->pParent->pParent->pLChild->color = BLACK;
                    z->pParent->pParent->color = RED;
                    z = z->pParent->pParent;
                }
                else {
                    if (z->pParent->pLChild == z) {
                        z = z->pParent;
                        RightRotate(z);
                    }

                    z->pParent->color = BLACK;
                    z->pParent->pParent->color = RED;
                    LeftRotate(z->pParent->pParent);
                }
            }
        }

        root->color = BLACK;
    }

    RBNode *DeleteNode(RBNode *node) {
        RBNode *p = NIL;
        if (node->pLChild==NIL || node->pRChild==NIL) p = node;
        else p = Successor(node);

        RBNode *pChild = NIL;
        if (p->pLChild != NIL)  pChild= p->pLChild;
        else pChild = p->pRChild;

        pChild->pParent = p->pParent;
        if (p->pParent == NIL) root = NIL;
        else if (p->pParent->pLChild == p) p->pParent->pLChild = pChild;
        else p->pParent->pRChild = pChild;

        if (p != node) {
            int tmp = node->key;
            node->key = p->key;
            p->key = tmp;
        }

        if (p->color == BLACK) DeleteFixup(pChild);

        if (pChild == NIL) pChild->pParent = NIL;

        return p;
    }

    void DeleteFixup(RBNode *x) {
        RBNode *w = NIL;
        while (x != root && x->color==BLACK) {
            if (x->pParent->pLChild == x) {
                w = x->pParent->pRChild;
                if (w->color == RED) {
                    w->color = BLACK;
                    x->pParent->color = RED;
                    LeftRotate(x->pParent);
                    w = x->pParent->pRChild;
                }

                if (w->pLChild->color==BLACK && w->pRChild->color==BLACK) {
                    w->color = RED;
                    x = x->pParent;
                }
                else {
                    if (w->pRChild->color == BLACK) {
                        w->pLChild->color = BLACK;
                        w->color = RED;
                        RightRotate(w);
                        w = w->pParent;
                    }

                    w->color = x->pParent->color;
                    x->pParent->color = BLACK;
                    w->pRChild->color = BLACK;
                    LeftRotate(x->pParent);
                    x = root;
                }

            }
            else {
                w = x->pParent->pLChild;
                if (w->color == RED) {
                    w->color = BLACK;
                    x->pParent->color = RED;
                    RightRotate(x->pParent);
                    w = x->pParent->pLChild;
                }

                if (w->pLChild->color==BLACK && w->pRChild->color==BLACK) {
                    w->color = RED;
                    x = x->pParent;
                }
                else {
                    if (w->pLChild == BLACK) {
                        w->pRChild->color = BLACK;
                        w->color = RED;
                        LeftRotate(w);
                        w = w->pParent;
                    }

                    w->color = x->pParent->color;
                    x->pParent->color = BLACK;
                    w->pLChild->color = BLACK;
                    RightRotate(x->pParent);
                    x = root;
                }
            }
        }
        x->color = BLACK;
    }

private:
    RBNode *NIL;
    RBNode *root;
};

int main() {

    int a[10] = {2, 5, 7, 8, 1, 9 ,10, 3, 4, 6};

    CRBTree btree;
    for (int i=0; i<10; ++i)
        btree.InsertKey(a[i]);

    printf("%d %d\n", btree.GetMinKey(), btree.GetMaxKey());

    btree.DeleteKey(1);
    btree.DeleteKey(10);
    printf("%d %d\n", btree.GetMinKey(), btree.GetMaxKey());

    btree.DeleteKey(5);
    if (btree.KeyInTree(5))
        puts("have found 5 in btree");
    else
        puts("can not find 5 in btree");

    if (btree.KeyInTree(4))
        puts("have found 4 in btree");
    else
        puts("can not find 4 in btree");

    return 0;
}