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

红黑树

简介

一直想写的一种数据结构，很是厉害的思想，插入，删除，查找，修改，都是logn的时间复杂度。
比AVL更强大的是，插入删除综合效率比AVL要优秀一点。
插入可能须要屡次变色，但旋转最多3次。删除可能须要屡次变色，但旋转最多3次。

# 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("==\n");
    } else {
        printf("!=\n");
    }

    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\n");
        else printf("OK!\n");
    }

    return 0;
}