红黑树实现(c/c++)

红黑树

简介

一直想写的一种数据结构，非常厉害的思想，插入，删除，查找，修改，都是(log_2 n)的时间复杂度。
比AVL更强大的是，插入删除综合效率比AVL要优秀一点。

性质

一颗红黑树是满足红黑性质的二叉搜索树:

1. 每个节点是红色或者黑色的。
2. 根节点是黑色的。
3. 每个叶节点(NIL)是黑色的
4. 如果一个节点是红色的，那么它的两个子节点都是黑色的。
5. 对于每个节点，从当前节点到其所有后代叶节点的简单路径上，黑节点的个数是相同的。

插入

插入肯定是插入到某叶节点的位置，颜色设为红色。所以可能违反性质4，按一定的规则修复。
插入修复可能需要多次变色，但旋转最多2次。
共分为3种情况(是父节点的左右孩子，对称共8种情况，以本节点是父节点的左孩子为例)：

1. 叔节点为红色。
2. 叔节点为黑色，且本节点为父节点的右孩子。
3. 叔节点为黑色，且本节点为父节点的左孩子。

2和3情况需要旋转。情况1可能变为1，2，3。但2只能变为3，3完成修复，所以最多2次旋转。

删除

删除一定是非叶节点，在二叉搜索树的删除方法上做了一定改动，并且按一定规则修复。
删除修复可能需要多次变色，但旋转最多3次。
共分为4种情况(是父节点的左右孩子，对称共8种情况，以本节点是父节点的左孩子为例)：

1. 兄弟节点为红色，且左右孩子节点都为黑色。
2. 兄弟节点为黑色，且左右孩子节点都为黑色。
3. 兄弟节点为黑色，且左孩子节点为红色，右孩子节点为黑色。
4. 兄弟节点为黑色，且左孩子节点为黑色，右孩子节点为红色。

实现

# include <cstdio>
# include <iostream>
using namespace std;

/**************************/

/*
红黑树的定义：
1.每个结点要么是红色，要么是黑色。
2.根结点是黑色的。
3.每个叶结点(NIL)是黑色的。
4.如果一个结点是红色的，那么它的两个子结点是黑色的。
5.每个结点到后代的叶结点的简单路径上的黑色结点个数相同。
*/

// 规定颜色
const int RED = 0;
const int BLACK = 1;

struct RBTreeNode {
    int key;
    int color;          // 颜色
    RBTreeNode * p;     // 父节点
    RBTreeNode * left;  // 左孩子
    RBTreeNode * right; // 右孩子
} * NIL, * root;


/// 初始化
void init() {
    NIL = new RBTreeNode;
    NIL->color = BLACK;
    NIL->p = NULL;
    NIL->left = NULL;
    NIL->right = NULL;
    root = NIL;
}


/// 新建节点
RBTreeNode * create_node(int key) {
    RBTreeNode * p = new RBTreeNode;
    p->key = key;
    p->p = NIL;
    p->left = NIL;
    p->right = NIL;
    p->color = RED;
    return p;
}


/// 根据键查询
RBTreeNode * search_node(int key) {
    RBTreeNode * x = root;
    while(x!=NIL && x->key != key) {
        if (key < x->key) x = x->left;
        else x = x->right;
    }
    return x;
}


/// 查找某子树最小结点
RBTreeNode * search_minimum(RBTreeNode * p) {
    if (p == NIL) return NIL;
    while(p->left != NIL) p = p->left;
    return p;
}


/// 查找某子树最大结点
RBTreeNode * search_maximum(RBTreeNode * p) {
    if (p == NIL) return NIL;
    while (p->right != NIL) p = p->right;
    return p;
}


/// 查询结点前驱结点(结点)
RBTreeNode * search_predecessor(RBTreeNode * p) {
    if (p == NIL) return NIL;
    if (p->left != NIL) {
        return search_maximum(p->left);  // 拥有左子树，后继一定是左子树的最大节点
    } else {
        RBTreeNode * y = p->p;
        while(y!=NIL && y->left==p) {  // 找到高层节点中以p所在的树为右子树的树的根节点，即是前驱节点
            p = y;
            y = y->p;
        }
        return y;
    }
}


/// 查找结点后继节点(结点)
RBTreeNode * search_successor(RBTreeNode * p) {
    if (p == NIL) return NIL;
    if (p->right != NIL) {
        return search_minimum(p->right);  // 拥有右子树，后继一定是右子树的最小节点
    } else {
        RBTreeNode * y = p->p;
        while(y!=NIL && y->right==p) {  // 找到高层节点中以p所在的树为左子树的树的根节点，即是后继节点
            p = y;
            y = y->p;
        }
        return y;
    }
}


/// 替换子树, u被v替换
void transplant_node(RBTreeNode * u, RBTreeNode * v) {
    if (u->p == NIL) {
        root = v;
    } else if (u->p->left == u) {
        u->p->left = v;
    } else u->p->right = v;
    if (v != NIL) {
        v->p = u->p;
    }
}


/// 结点左旋(x, y不为NIL)
void left_rotate(RBTreeNode * x) {
    RBTreeNode * y = x->right;
    transplant_node(x, y);
    RBTreeNode * z = y->left;
    x->p = y;
    y->left = x;
    if (z != NIL) z->p = x;
    x->right = z;
}


/// 结点右旋(x, y不为NIL)
void right_rotate(RBTreeNode * x) {
    RBTreeNode * y = x->left;
    transplant_node(x, y);
    RBTreeNode * z = y->right;
    x->p = y;
    y->right = x;
    if (z != NIL) z->p = x;
    x->left = z;
}


/// 插入结点调整
void insert_node_fixup(RBTreeNode * x) {
     while (x->p->color == RED) {
        RBTreeNode * y = x->p;
        if (y->p->left == y) {  // 位于爷结点的左子树
            RBTreeNode * z = y->p->right;
            if (z->color == RED) {  // case1: 叔结点是红色
                z->color = BLACK;
                y->color = BLACK;
                y->p->color = RED;
                x = y->p;
                continue;
            }
            if (y->right == x) {  // case2: 叔结点是黑色，是父结点的右孩子
                x = x->p;
                left_rotate(y);
            }
            x->p->color = BLACK;  // case3: 叔结点是黑色，是父结点的左孩子
            x->p->p->color = RED;
            right_rotate(x->p->p);
        } else {  // 位于爷结点的右子树
            RBTreeNode * z = y->p->left;
            if (z->color == RED) {
                z->color = BLACK;
                y->color = BLACK;
                y->p->color = RED;
                x = y->p;
                continue;
            }
            if (y->left == x) {
                x = x->p;
                right_rotate(y);
            }
            x->p->color = BLACK;
            x->p->p->color = RED;
            left_rotate(x->p->p);
        }
    }
    root->color = BLACK;
}


/// 插入结点(结点)
void insert_node(RBTreeNode * z) {
    RBTreeNode * x = root;
    RBTreeNode * y = NIL;
    while (x!=NIL) {
        y = x;
        if (z->key < x->key) x = x->left;
        else x = x->right;
    }
    z->p = y;
    if (y == NIL)
        root = z;
    else if (z->key < y->key)
        y->left = z;
    else
        y->right = z;
    insert_node_fixup(z);
}


/// 调整删除结点
void delete_node_fixup(RBTreeNode * x) {
    while(x != root && x->color == BLACK) {
        if (x->p->left == x) {
            RBTreeNode * w = x->p->right;
            if (w->color == RED) {  // case1: 兄弟结点是红色
                x->p->color = RED;
                w->color = BLACK;
                left_rotate(x->p);
            }
            if (w->left->color == BLACK && w->right->color == BLACK) {  // case2: 兄弟结点是黑色，并且双亲为黑色
                w->color = RED;
                x = x->p;
                continue;
            }
            if (w->right->color != RED) {  // case3: 兄弟结点是黑色，左孩子为红色
                w->left->color = BLACK;
                w->color = RED;
                right_rotate(w);
            }
            // case4: 兄弟结点是黑色，右孩子是红色
            w->color = x->p->color;
            w->right->color = BLACK;
            x->p->color = BLACK;
            left_rotate(x->p);
            x = root;
        } else {
            RBTreeNode * w = x->p->left;
            if (w->color == RED) {  // case1: 兄弟结点是红色
                x->p->color = RED;
                w->color = BLACK;
                right_rotate(x->p);
            }
            if (w->right->color == BLACK && w->left->color == BLACK) {  // case2: 兄弟结点是黑色，并且双亲为黑色
                w->color = RED;
                x = x->p;
                continue;
            }
            if (w->left->color != RED) {  // case3: 兄弟结点是黑色，左孩子为红色
                w->right->color = BLACK;
                w->color = RED;
                left_rotate(w);
            }
            // case4: 兄弟结点是黑色，右孩子是红色
            w->color = x->p->color;
            w->left->color = BLACK;
            x->p->color = BLACK;
            right_rotate(x->p);
            x = root;
        }
    }
    x->color = BLACK;
}


/// 删除结点(结点)
void delete_node(RBTreeNode * z) {
    RBTreeNode * x;  // 记录被删除的结点在树中所处的位置
    RBTreeNode * y = z;  // 记录实际被删除的结点
    int y_origin_color = y->color;
    if (z->left == NIL) {
        x = z->right;
        transplant_node(z, z->right);
    } else if (z->right == NIL) {
        x = z->left;
        transplant_node(z, z->left);
    } else {  // 左右孩子都存在的情况
        y = search_minimum(z->right);  // 找后继节点
        y_origin_color = y->color;
        x = y->right;
        if (y != x->right) {  // 如果后继不是右孩子，需要变形。将后继节点提为右子树的根节点
            transplant_node(y, y->right);  // 后继节点的左孩子一定不存在，右孩子取代后继节点
            y->right = z->right;
            y->right->p = y;
        }
        // 后继就是右孩子
        transplant_node(z, y);
        y->left = z->left;  // 替换后还需要修改与左子树的父子关系与颜色
        z->left->p = y;
        y->color = z->color;
    }
    delete z;
    if (y_origin_color == BLACK) delete_node_fixup(x);
}


/** --- */

bool insert_node(int key) {
    RBTreeNode * node = search_node(key);
    if (node != NIL) return false;
    node = create_node(key);
    insert_node(node);
    return true;
}

bool delete_node(int key) {
    RBTreeNode * node = search_node(key);
    if (node == NIL) return false;
    delete_node(node);
    return true;
}


int main() {
    init();

    RBTreeNode * x = NIL;
    if (x == NIL){
        printf("==
");
    } else {
        printf("!=
");
    }

    insert_node(1);
//    insert_node(3);
//    insert_node(5);
//    insert_node(7);
//    insert_node(9);
//
//    insert_node(2);
//    insert_node(4);
//    insert_node(6);
//    insert_node(8);
//    insert_node(10);
//
//    delete_node(3);
//    delete_node(7);
//    delete_node(6);
//    delete_node(1);

    while (1) {
        int k;
        scanf("%d", &k);
        RBTreeNode * p = search_node(k);
        if (p == NIL) printf("NIL
");
        else printf("OK!
");
    }

    return 0;
}