Red-black tree implementation (c / c ++)

Red-black tree

Brief introduction

Always wanted to write a data structure, a very powerful idea, insert, delete, search, modify, all logn time complexity.
AVL is more powerful than the insertion deletion overall efficiency than AVL to excellent point.
Insert may be necessary to change color, but the rotation up to 3 times. Delete may take several color, but the rotation up to 3 times.

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