算法导论第三版--红黑树

红黑树是满足如下条件的二叉搜索树。

1、每个节点要么是红色的，要么是黑色的。

2、根节点是黑色。

3、每个叶子节点是黑色的。

4、如果一个节点是红色的，那么它的2个子节点是黑色的。

5、对每个节点，从它到它的后代叶子节点的简单路径上，均包含相同数目的黑色节点。

练习完这些代码，我感觉我把stl的rbtree抄了一遍。。。

enum RB {
    RED,
    BLACK
};

template<typename T>
struct rb_node_cls {
    typedef rb_node_cls<T> __node;
    __node* parent;
    __node* left;
    __node* right;
    T key;
    RB color;
    void *data;
};

template<typename T> using rbt_node = rb_node_cls<T>;

template<typename T>
struct rb_tree_cls {
   rbt_node<T>* root;
};

template<typename T> using rb_tree = rb_tree_cls<T>;

template<typename T>
void left_rotate(rb_tree<T>& tree, rbt_node<T>* x) {
    rbt_node<T>* y = x->right;

    x->right = y->left;
    if (y->left != nullptr) {
        y->left->parent = x;
    }

    y->parent = x->parent;
    if(x->parent == nullptr) {
        tree.root = y;
    } else if(x == x->parent->left) {
        x->parent->left = y;
    } else {
        x->parent->right = y;
    }

    y->left = x;
    x->parent = y;
}

template<typename T>
void right_rotate(rb_tree<T>& tree, rbt_node<T>* x) {
    rbt_node<T>* y = x->left;

    y->left = x->right;

    if(x->right != nullptr) {
        x->right->parent = y;
    }

    x->parent = y->parent;
    if(y->parent == nullptr) {
        tree.root = y;
    } else if (y == y->parent->left) {
        y->parent->left = x;
    } else {
        x->parent->right = x;
    }

    x->right = y;
    y->parent = x;
}

template<typename T>
void rb_insert_fixup(rb_tree<T>& tree, rbt_node<T>* z) {
    while(z->parent->color == RED) {
        if(z->parent == z->parent->parent->left) {
            rbt_node<T>* y = z->parent->parent->right;
            if(y->color == RED){
                z->parent->color = BLACK;
                y->color = BLACK;
                z->parent->parent = RED;
                z = z->parent->parent;
            } else if(z == z->parent->right) {
                z = z->parent;
                left_rotate(tree, z);
            }
            z->parent->color = BLACK;
            z->parent->parent->color = RED;
            right_rotate(tree, z->parent->parent);
        } else {
            rbt_node<T>* y = z->parent->parent->left;
            if(y->color == RED) {
                z->parent->color = BLACK;
                y->color = BLACK;
                z->parent->parent = RED;
                z = z->parent->parent;
            } else if(z == z->parent->left) {
                z = z->parent;
                right_rotate(tree, z);
            }
            z->parent->color = BLACK;
            z->parent->parent->color = RED;
            left_rotate(tree, z->parent->parent);
        }
    }
    tree.root->color = BLACK;
}

template<typename T>
void rb_insert(rb_tree<T>& tree, rbt_node<T>* z) {
    rbt_node<T>* y = nullptr;
    rbt_node<T>* x = tree.root;

    while(x != nullptr) {
        y = x;
        x = (z->key < x->key) ? x->left : x->right;
    }

    z->parent = y;
    if(y == nullptr){
        tree.root = z;
    } else if(z->key < y->key) {
        y->left = z;
    } else {
        y->right = z;
    }

    z->left = nullptr;
    z->right = nullptr;
    z->color = RED;
    rb_insert_fixup(tree, z);
}

template<typename T>
void rb_transplant(rb_tree<T>& tree, rbt_node<T>* u, rbt_node<T>* v) {
    if(u->parent == nullptr) {
        tree.root = v;
    } else if (u == u->parent->left) {
        u->parent->left = v;
    } else {
        u->parent->right = v;
    }
    if(v != nullptr)
        v->parent = u->parent;
}

template<typename T>
rbt_node<T>* tree_minimum(rbt_node<T>* node) {
    while(node->left != nullptr)
        node = node->left;
    return node;
}

template<typename T>
void rb_delete_fixup(rb_tree<T>& tree, rbt_node<T>* x) {
    while(x != tree.root && x->color == BLACK) {
        if(x == x->parent->left) {
            rbt_node<T>* w = x->parent->right;
            if(w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                left_rotate(tree, x->parent);
                w = x->parent->right;
            }

            if(w->left->color == BLACK && w->right->color == BLACK) {
                w->color = RED;
                x = x->parent;
            } else if(w->right->color == BLACK) {
                w->left->color = BLACK;
                w->color = RED;
                right_rotate(tree, w);
                w = x->parent->right;
            }

            w->color = x->parent->color;
            x->parent->color = BLACK;
            w->right->color = BLACK;
            left_rotate(tree, x->parent);
            x = tree.root;
        } else {
            rbt_node<T>* w = x->parent->left;
            if(w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                right_rotate(tree, x->parent);
                w = x->parent->left;
            }

            if(w->right->color == BLACK && w->left->color == BLACK) {
                w->color = RED;
                x = x->parent;
            } else if (w->left->color == BLACK) {
                w->right->color = BLACK;
                w->color = RED;
                left_rotate(tree, w);
                w = x->parent->left;
            }

            w->color = x->parent->color;
            x->parent->color = BLACK;
            x->left->color = BLACK;
            right_rotate(tree, x->parent);
            x = tree.root;
        }
    }
    x->color = BLACK;
}

template<typename T>
void rb_delete(rb_tree<T>& tree, rbt_node<T>* z) {
    rbt_node<T>* y = z;
    rbt_node<T>* x = nullptr;
    RB y_original_color = y->color;

    if(z->left == nullptr) {
        x = z->right;
        rb_transplant(tree, z, z->right);
    } else if(z->right == nullptr) {
        x = z->left;
        rb_transplant(tree, z, z->left);
    } else {
        y = tree_minimum(z->right);
        y_original_color = y->color;
        x = y->right;
        if(y->parent == z){
            x->parent = y;
        } else {
            rb_transplant(tree, y, y->right);
            y->right = z->right;
            y->right->parent = y;
        }
        rb_transplant(tree, z, y);
        y->left = z->left;
        y->left->parent = y;
        y->color = z->color;
    }
    if(y_original_color == BLACK){
        rb_delete_fixup(tree, x);
    }
}