红黑树

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

#define NIL (&__NIL)

typedef struct Node {
        struct Node *lchild, *rchild;
        int color, val;
} Node;

Node __NIL;

__attribute__((constructor))
void init_NIL() {
        NIL->lchild = NIL->rchild = NIL;
        NIL->color = 1, NIL->val = -1;
        return ;
}

Node *getNewNode(int a) {
        Node *p = (Node *)malloc(sizeof(Node));
        p->color = 0, p->val = a;
        p->lchild = p->rchild = NIL;
        return p;
}

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

int has_red_child(Node *root) {
        return !root->lchild->color || !root->rchild->color;
}

Node *left_rotate(Node *root) {
        printf("left%d\n", root->val);
        Node *temp = root->rchild;
        root->rchild = temp->lchild;
        temp->lchild = root;
        return temp;
}

Node *right_rotate(Node *root) {
        printf("right%d\n", root->val);
        Node *temp = root->lchild;
        root->lchild = temp->rchild;
        temp->rchild = root;
        return temp;
}

Node *erase_maintain(Node *root) {
        if (root->lchild->color != 2 && root->rchild->color != 2) return root;
        if (has_red_child(root)) {
                root->color = 0;
                root->lchild->color && (root = left_rotate(root), root->color = 1, root->lchild = erase_maintain(root->lchild)) || 
(root = right_rotate(root), root->color = 1, root->rchild = erase_maintain(root->rchild));
                return root;
        } else {
                if((root->lchild->color == 2 && !has_red_child(root->rchild)) ||
                (root->rchild->color == 2 && !has_red_child(root->lchild)))
                        root->color++, root->lchild->color--, root->rchild->color--;
                else {
                        if (root->lchild->color == 2) {
                                root->lchild->color = 1;
                                root->rchild->lchild->color || (root->rchild->color = 0, root->rchild = right_rotate(root->rchild), root->rchild->color = 1);
                                root->rchild->color = root->color;
                                root = left_rotate(root);
                                root->lchild->color = root->rchild->color = 1;
                        } else {
                                root->rchild->color = 1;
                                root->lchild->rchild->color || (root->lchild->color = 0, root->lchild = left_rotate(root->lchild), root->lchild->color = 1);
                                root->lchild->color = root->color;
                                root = right_rotate(root);
                                root->rchild->color = root->lchild->color = 1;
                        }
                }
        }
        return root;
}

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

Node *__erase(Node *root, int key) {
        if (root == NIL) return NIL;
        root->val > key && (root->lchild = __erase(root->lchild, key));
        root->val < key && (root->rchild = __erase(root->rchild, key));
        if (root->val == key) {
                if (root->rchild == NIL || root->lchild == NIL) {
                        Node *temp = root->rchild == NIL ? root->lchild : root->rchild;
                        temp->color += root->color;
                        free(root);
                        return (temp);
                }
                Node *temp = predecessor(root);
                root->val = temp->val;
                root->lchild = __erase(root->lchild, root->val);
        }
        return erase_maintain(root);
}

Node *erase(Node *root, int key) {
        if (root == NIL) return root;
        printf("\nerase: %d\n", key);
        root = __erase(root, key);
        root->color = 1;
        return root;
}

Node *insert_maintain(Node *root) {
        if (root == NIL || !has_red_child(root))  return root;
        int a;
        if (root->lchild->color == 0 && root->rchild->color == 0) {
                if (!has_red_child(root->lchild) && !has_red_child(root->rchild)) return root;
                if (has_red_child(root->lchild) || has_red_child(root->rchild)) {
                        root->lchild->color = root->rchild->color = 1;
                        root->color = 0;
                }
                return root;
        }
        if (!root->lchild->color) {
                if (!has_red_child(root->lchild)) return root;
                !root->lchild->rchild->color && (root->lchild = left_rotate(root->lchild));
                root = right_rotate(root);
        }
        if (!root->rchild->color) {
                if (!has_red_child(root->rchild)) return root;
                !root->rchild->lchild->color && (root->rchild = right_rotate(root->rchild));
                root = left_rotate(root);
        }
        root->color = 0, root->lchild->color = root->rchild->color = 1;
        return root;
}

Node *__insert(Node *root, int a) {
        if (root == NIL) return getNewNode(a);
        if (root->val == a) return root;
        root->val < a && (root->rchild = __insert(root->rchild, a));
        a < root->val && (root->lchild = __insert(root->lchild, a));
        return insert_maintain(root);
}

Node *insert(Node *root, int a) {
        printf("insert: %d\n");
        if (!root) return getNewNode(a);
        root = __insert(root, a);
        root->color = 1;
        return root;
}



void output(Node *root) {
        if (root == NIL) return ;
        printf("%d, color = %d | (%d, %d)\n", root->val, root->color, root->lchild->val, root->rchild->val);
        output(root->lchild);
        output(root->rchild);
        return ;
}

int main() {
        int a;
        srand(time(0));
        scanf("%d", &a);
        Node *root;
        for (int i = 0; i < a; i++) root = insert(root, rand() % 100);
        output(root);
        while (~scanf("%d", &a)) root = erase(root, a), output(root);
        return 0;
}

posted @ 2021-09-11 05:17  代码附体  阅读(107)  评论(0编辑  收藏  举报