紅黑樹實現(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("==\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;
}