手撕红黑树

来源：互联网发布：数据库长整型编辑：程序博客网时间：2024/06/07 14:52
#include <stdio.h>#include <stdlib.h>typedef struct RBTNode {    int key;    int black;    struct RBTNode *lchild, *rchild;    struct RBTNode *father;} RBTNode;RBTNode *NIL;#define lson node->lchild#define rson node->rchild#define swap(x,y) x=x+y;y=x-y;x=x-yvoid init_NIL() {    NIL = (RBTNode *)malloc(sizeof(RBTNode));    NIL->father = NULL;    NIL->lchild = NIL->rchild = NIL;    NIL->key = 0;    NIL->black = 1;}RBTNode *init(int key) {    RBTNode *p = (RBTNode *)malloc(sizeof(RBTNode));    p->father = NULL;    p->lchild = p->rchild = NIL;    p->key = key;    p->black = 0;    return p;}void clear(RBTNode *node) {    if (node == NULL || node == NIL) return ;    clear(lson);    clear(rson);    free(node);    return ;}void clear_NIL() {    free(NIL);    NIL = NULL;}int has_red_child(RBTNode *node) {    return lson->black == 0 || rson->black == 0;}RBTNode *left_rotate(RBTNode *node) {    RBTNode *temp = rson;    temp->father = node->father;    node->father = temp;    temp->lchild->father = node;    rson = temp->lchild;    temp->lchild = node;    return temp;}RBTNode *right_rotate(RBTNode *node) {    RBTNode *temp = lson;    temp->father = node->father;    node->father = temp;    temp->rchild->father = node;    lson = temp->rchild;    temp->rchild = node;    return temp;}int judge_insert(RBTNode*node){    if(lson->black+lson->lchild->black==0) return 1;    if(lson->black+lson->rchild->black==0) return 2;    if(rson->black+rson->lchild->black==0) return 3;    if(rson->black+rson->rchild->black==0) return 4;    return 0;}RBTNode *insert_adjust(RBTNode *node) {    if (node->black == 0) return node;    int judge_result=judge_insert(node);    if(judge_result) {    if (lson->black == 0 && rson->black == 0) {        lson->black = rson->black = 1;        node->black = 0;        return node;    }    if (judge_result==1) {        node = right_rotate(node);        lson->black = 1;    } else if (judge_result==2) {        lson = left_rotate(lson);        node = right_rotate(node);        lson->black = 1;    } else if (judge_result == 4) {        node = left_rotate(node);        rson->black = 1;    } else if (judge_result == 3) {        rson = right_rotate(rson);        node = left_rotate(node);        rson->black = 1;    }    }    return node;}RBTNode *insert(RBTNode *node, int key) {    if (node == NULL || node == NIL) {        RBTNode *temp = init(key);        temp->black = (node == NULL);        return temp;    }    if (node->key == key) return node;    if (node->key < key) {        rson = insert(rson, key);        rson->father = node;        node = insert_adjust(node);    } else {        lson = insert(lson, key);        lson->father = node;        node = insert_adjust(node);    }    if (node->father == NULL) {        node->black = 1;    }    return node;}RBTNode *predecessor(RBTNode *node) {    RBTNode *temp = lson;    while (temp->rchild != NIL) temp = temp->rchild;    return temp;}RBTNode *erase_adjust(RBTNode *node) {    if (lson->black != 2 && rson->black != 2) return node;    if (lson->black == 2 && rson->black == 0) {        node = left_rotate(node);        node->black = 1;        lson->black = 0;        lson = erase_adjust(lson);    } else if (rson->black == 2 && lson->black == 0) {        node = right_rotate(node);        node->black = 1;        rson->black = 0;        rson = erase_adjust(rson);    } else if (lson->black == 2 && rson->black == 1) {        if (!has_red_child(rson)) {            rson->black = 0;            lson->black = 1;            node->black++;            return node;        } else if (rson->rchild->black != 0) {            rson = right_rotate(rson);            rson->black = 1;            rson->rchild->black = 0;        }        node = left_rotate(node);        node->black = lson->black;        lson->black = rson->black = 1;        lson->lchild->black = 1;    } else if (rson->black == 2 && lson->black == 1) {        if (!has_red_child(lson)) {            lson->black = 0;            rson->black = 1;            node->black++;            return node;        } else if (lson->lchild->black != 0) {            lson = left_rotate(lson);            lson->black = 1;            lson->lchild->black = 0;        }        node = right_rotate(node);        node->black = rson->black;        lson->black = rson->black = 1;        rson->rchild->black = 1;    }    return node;}RBTNode *erase(RBTNode *node, int key) {    if (node == NULL || node == NIL) {        return node;    }    if (key > node->key) {        rson = erase(rson, key);        node = erase_adjust(node);    } else if (key < node->key) {        lson = erase(lson, key);        node = erase_adjust(node);    } else {        if (lson == NIL && rson == NIL) {            NIL->black += node->black;            free(node);            return NIL;        } else if (lson == NIL || rson == NIL) {            RBTNode *temp = (lson == NIL ? rson : lson);            temp->black += node->black;            temp->father = node->father;            free(node);            return temp;        } else {            RBTNode *temp = predecessor(node);            swap(node->key,temp->key);            lson = erase(lson, key);            node = erase_adjust(node);        }    }    if (node->father == NULL) {        node->black = 1;    }    return node;}void inorder(RBTNode *node) {    if (node == NULL || node == NIL) return ;    inorder(lson);    printf("%d %d %d %d\n", node->key, node->black, lson->key, rson->key);    inorder(rson);    return ;}int main() {    init_NIL();    int opr, n;    RBTNode *root = NULL;    while (scanf("%d%d", &opr, &n) != EOF) {        switch (opr) {            case 1:                root = insert(root, n);                break;            case 2:                root = erase(root, n);                break;            case 3:               inorder(root);        }    }    clear(root);    clear_NIL();    return 0;}
阅读全文
0 0