手撕红黑树

#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-y

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