《高级数据结构》-红黑树

思维导图

代码演示

效果

代码

/*************************************************************************
	> File Name: red_black_insert.cpp
	> Mail: 1136984246@qq.com
 ************************************************************************/

#include <stdlib.h>
#include <stdio.h>


typedef struct Node {
    int key;
    int color; // 0 : red, 1 black, 2 double black
    struct Node *lchild, *rchild;
} Node;

Node __NIL;
#define NIL (&__NIL) 

__attribute((constructor))
void init_NIL() {
    NIL->key = 0; 
    NIL->color = 1; // NIL作为叶子节点，为黑色
    NIL->lchild = NIL->rchild = NIL;
}

Node *getNewNode(int key) {
    Node *p = (Node *)malloc(sizeof(Node));
    p->key = key;
    p->color = 0; // 对于插入的话，默认插入红色
    p->lchild = p->rchild = NIL;

}

// 有无红孩子，至少有一个红孩子
int has_red_child(Node *root) {
    return root->lchild->color == 0 || root->rchild->color == 0; 
}

// 左旋：同AVL树
Node *left_rotate(Node *root) {
    Node *temp = root->rchild;
    root->rchild = temp->lchild;
    temp->lchild = root;
    return temp;
}

// 右旋
Node *right_rotate(Node *root) {
    Node *temp = root->lchild;
    root->lchild = temp->rchild;
    temp->rchild = root;
    return temp;
}

// 插入调整 
Node *insert_maintain(Node *root) {
    int flag = 0;
    if (!has_red_child(root)) return root; // 没有红孩子，不需要调整
    // 情况一：不管冲突不冲突[只要是黑红红，而不用管红下面有没有红]，直接红黑黑
    if (root->lchild->color == 0 && root->rchild->color == 0) goto insert_end;
    // 情况二：需判断冲突发生在左子树，还是发生在右子树，需要使用flag进行标识
    if (root->lchild->color == 0 && has_red_child(root->lchild)) flag = 1; // 左子树发生冲突
    if (root->rchild->color == 0 && has_red_child(root->rchild)) flag = 2; // 右子树发生冲突
    // 无冲突，不调整
    if (flag == 0) return root;
    // 有冲突：旋转（同AVL) + 染色
    if (flag == 1) {
        // LR, LL
        if (root->lchild->rchild->color == 0) {
            root->lchild = left_rotate(root->lchild);
        }
        root = right_rotate(root);
    } else {
        // RL, RR
        if (root->rchild->lchild->color == 0) {
            root->rchild = right_rotate(root->rchild);
        }
        root = left_rotate(root);
    }
// 染色：采用红色上浮方式，也可才用红色下沉（此时需区分情况一还是情况二）
insert_end:
    root->color = 0;
    root->lchild->color = root->rchild->color = 1;
    return root;
}

// 插入：封装一层
Node *__insert(Node *root, int key) {
    if (root == NIL) return getNewNode(key); // 此时插入节点
    if (root->key == key) return root;
    if (key < root->key) {
        root->lchild = __insert(root->lchild, key);
    } else {
        root->rchild = __insert(root->rchild, key);
    }
    return insert_maintain(root); // 回溯的过程中进行调整
}

// 插入 + 染色[必须的，保证2nd条件：黑节点的来源]
Node *insert(Node *root, int key) {
    root = __insert(root, key);
    root->color = 1; // 手头染黑，最上面的根节点
    return root;
}


Node *predecessor(Node *root) {
    Node *temp = root->lchild;
    while (temp->rchild != NIL) temp = temp->rchild;
    return temp;
}

// 删除调整
Node *erase_maintain(Node *root) {
    // 没有双重黑子节点，不需要调整
    if (root->lchild->color != 2 && root->rchild->color != 2) return root;
    // 有双黑子节点，需要分情况来讨论
    // [特殊情况]：双重黑的兄弟节点是红色 -> 调整成黑色
    if (has_red_child(root)) {
        int flag = 0; // 记录旋转方向
        root->color = 0; // 原根节点改成红色
        if (root->lchild->color == 0) {
            // 红色在左边/右边 -> 右旋/左旋
            root = right_rotate(root); flag = 1;
        } else {
            root = left_rotate(root); flag = 2;
        }
        root->color = 1; // 新根节点改成黑色
        // 右旋/左旋 -> 双黑节点的父节点(原根节点)分别再右/左子树中, 转化成情况一、二、三
        // 对于原根节点左删除调整
        if (flag == 1) root->rchild = erase_maintain(root->lchild);
        else root->lchild = erase_maintain(root->rchild);
        return root;
    } 
    // [三种情况] : 双黑子节点的兄弟节点为黑色节点
    // [情况一]：双黑子节点的兄弟节点没有红孩子
    if ((root->lchild->color == 2 && !has_red_child(root->rchild)) ||
        (root->rchild->color == 2 && !has_red_child(root->lchild))) {
            // 左右子节点减少一重黑（解决了双黑），根节点加一重黑
            root->lchild->color -= 1;
            root->rchild->color -= 1;
            root->color += 1;
            return root;
    }
    // [情况二、三] 此时兄弟节点一定有红孩子
    if (root->lchild->color == 2) {
    // RL, RR
        root->lchild->color -= 1; // 解决双黑 [顺序无讲究]
        // 非红：可能是黑，可能是双重黑（NIL是公用的） 不用：root->rchild->rchild->color == 1 
        if (root->rchild->rchild->color != 0) {
            // RL
            // 原兄弟节点染成红色, 小右旋，新兄弟节点染成黑色
            root->rchild->color = 0;
            root->rchild = right_rotate(root->rchild);
            root->rchild->color = 1;
        } 
        root = left_rotate(root); // 大左旋（RR，RL均有）
        root->color = root->lchild->color; // 新根节点的颜色等于原根节点的颜色
    } else {
    // LR, LL
        root->rchild->color -= 1;
        if (root->lchild->lchild->color != 0) {
            root->lchild->color = 0;
            root->lchild = left_rotate(root->lchild);
            root->lchild->color = 1;
        }
        root = right_rotate(root);
        root->color = root->rchild->color;
    }
    root->lchild->color = root->rchild->color = 1; // 两个新的子节点强制变成黑色

    return root;
}


Node *__erase(Node *root, int key) {
    if (root == NIL) return NIL;
    if (key < root->key) {
        root->lchild = __erase(root->lchild, key);
    } else if (key > root->key) {
        root->rchild = __erase(root->rchild, key);
    } else {
        if (root->lchild == NIL || root->rchild == NIL) {
            // 度为0或1 [度为0的节点删除 -> 双重黑]
            Node *temp = root->lchild != NIL ?root->lchild : root->rchild;
            temp->color += root->color; // 可处理2种情况
            free(root);
            return temp;
        } else {
            // 度为2 转化 度为0或1
            Node *temp =  predecessor(root);
            root->key = temp->key;
            root->lchild = __erase(root->lchild, temp->key);
        }
    }
    return erase_maintain(root);
} 

Node *erase(Node *root, int key) {
    root = __erase(root, key);
    root->color = 1; // 强制染黑头
    return root;
}

void clear(Node *root) {
    if (root == NIL) return;
    clear(root->lchild);
    clear(root->rchild);
    free(root);
    return;
}

// 根节点的颜色 | 根节点的值，左孩子的值，右孩子的值
void print(Node *root) {
    printf("%d| %d, %d, %d\n",
        root->color, root->key,
        root->lchild->key,
        root->rchild->key
    );
}

void output(Node *root) {
    if (root == NIL) return;
    print(root);
    output(root->lchild);
    output(root->rchild);
}

int main() {
    int op, val;
    Node *root = NIL;
    while (~scanf("%d%d", &op, &val)) {
        switch (op) {
            case 1: root = insert(root, val); printf("insert %d\n", val); break;
            case 2: root = erase(root, val);  printf("erase %d\n", val); break;
        } 
        output(root);
        printf("-------------------\n");
    }
    clear(root);
    return 0;
}

参考文献

《高级数据结构》-红黑树