性质

  • 节点是红色或黑色
  • 根节点是黑色
  • 所有叶子都是黑色。(叶子是NUIL节点)
  • 每个红色节点的两个子节点都是黑色。(从每个叶子到根的所有路径上不能有两个连续的红色节点)
  • 从任一节点到其每个叶子的所有路径都包含相同数目的黑色节点
    • 如果一个节点的一个子节点是黑色,那么这个节点一定有2个节点。

插入情景
图片说明

删除情景
图片说明

图片说明

Code

enum Color{red,black};
template<typename T>
struct BiTNode{
    BiTNode *left;
    BiTNode *right;
    BiTNode *parent;
    Color color;
    T val;
    int son;
    BiTNode():son(0){}
    BiTNode(T val):val(val),son(1){
        color=red;
        parent=NULL;
        left=NULL;
        right=NULL;
    }
    BiTNode *grandpa(){
        if(parent==NULL)return NULL;
        return parent->parent;
    }
    BiTNode* sibling(){
        if(parent->left==this)return parent->right;
        return parent->left;
    }
    BiTNode* uncle(){
        if(grandpa()==NULL)return NULL;
        if(grandpa()->left==parent)return grandpa()->right;
        return grandpa()->left;
    }
};
template<typename T,typename node=BiTNode<T>>
class rb_tree{
public:
    rb_tree():n(0),root(NULL){
        NIL=new node();
        NIL->color=black;
    }
    ~rb_tree(){clear(root);delete NIL;}
    rb_tree(T* bg,T* ed):n(0),root(NULL){
        NIL=new node();
        NIL->color=black;
        while(bg!=ed)insert(*bg++);
    }
    bool empty()const{return n==0;}
    int size()const{return n;}
    void insert(const T& key){insert(root,key);++n;}
    void erase(const T& key){remove(root,key);--n;}
    bool find(const T& key)const{return search(root,key)!=NULL;}
    void clear(){clear(root);n=0;}
    T& MIN(){return MIN(root)->val;}
    T& MAX(){return MAX(root)->val;}
    T kth(int k){
        if(k>n)return -1;
        node *p=root;
        while(p!=NIL){
            if(p->left!=NIL){
                if(k<=p->left->son)p=p->left;
                else{
                    k-=p->left->son;
                    if(k==1)return p->val;
                    k--;
                    p=p->right;
                }
            }else{
                if(k==1)return p->val;
                k--;
                p=p->right;
            }
        }
    }
private:
    #define is_red(p) ((p)->color==red)
    #define is_black(p) ((p)->color==black)
    #define set_red(p) ((p)->color=red)
    #define set_black(p) ((p)->color=black)
    void clear(node* &rt)
    {
        if(rt==NIL)return;
        clear(rt->left);
        clear(rt->right);
        delete rt;
    }
    node* search(node* rt,const T& key)
    {
        if(rt==NULL)return NULL;
        node* p=rt;
        while(p!=NIL)
        {
            if(key==p->val)return p;
            if(key<p->val)p=p->left;
            else p=p->right;
        }
        return NULL;
    }
    node* MIN(node* rt)
    {
        node* p=rt;
        while(p->left!=NIL)p=p->left;
        return p;
    }
    node* MAX(node* rt)
    {
        node* p=rt;
        while(p->right!=NIL)p=p->right;
        return p;
    }
    node* suf(node* rt)
    {
        if(rt->right!=NIL)return MIN(rt->right);
        node* p=rt->parent;
        while(p&&rt==p->right)rt=p,p=p->parent;
        return p;
    }
    node* pre(node* rt)
    {
        if(rt->left!=NIL)return MAX(rt->left);
        node* p=rt->parent;
        while(p&&rt==p->left)rt=p,p=p->parent;
        return p;
    }
    void insert(node* &rt,const T& key){
        if(rt==NULL){
            rt=new node(key);
            rt->left=rt->right=NIL;
            set_black(rt);
            return;
        }
        push(rt,key);
    }
    int num(node* &p){return p==NIL?0:p->son;}
    void up_son(node* &p){p->son=num(p->left)+num(p->right)+1;}
    bool lson(node* &p){return p->parent->left==p;}
    bool rson(node* &p){return p->parent->right==p;}
    void push(node* rt,const T& key){
        if(key<=rt->val){
            if(rt->left!=NIL)push(rt->left,key);
            else{
                node* p=new node(key);
                rt->left=p;
                p->parent=rt;
                p->left=p->right=NIL;
                up_insert(p);
            }
        }else if(key>rt->val){
            if(rt->right!=NIL)push(rt->right,key);
            else{
                node* p=new node(key);
                rt->right=p;
                p->parent=rt;
                p->left=p->right=NIL;
                up_insert(p);
            }
        }
        up_son(rt);
    }
    void up_insert(node *p){
        if(p->parent==NULL){
            root=p;
            set_black(p);
            return;
        }
        if(is_black(p->parent))return;
        if(is_red(p->uncle())){
            set_black(p->parent);
            set_black(p->uncle());
            set_red(p->grandpa());
            up_insert(p->grandpa());return;
        }
        if(lson(p)&&rson(p->parent)){
            rotate_right(p);
            p=p->right;
        }
        else if(rson(p)&&lson(p->parent)){
            rotate_left(p);
            p=p->left;
        }
        if(lson(p)&&lson(p->parent)){
            set_black(p->parent);
            set_red(p->grandpa());
            rotate_right(p->parent);
        }
        else if(rson(p)&&rson(p->parent)){
            set_black(p->parent);
            set_red(p->grandpa());
            rotate_left(p->parent);
        }
    }
    void rotate_left(node *p){
        node *grandpa=p->grandpa();
        node *parent=p->parent;
        node *left=p->left;
        parent->right=left;
        if(left!=NIL)left->parent=parent;
        p->left=parent;
        parent->parent=p;
        p->parent=grandpa;
        if(grandpa==NULL)root=p;
        else{
            if(grandpa->left==parent)
            grandpa->left=p;
            else grandpa->right=p;
        }
        up_son(parent);up_son(p);
    }
    void rotate_right(node *p){
        node *grandpa=p->grandpa();
        node *parent=p->parent;
        node *right=p->right;
        parent->left=right;
        if(right!=NIL)right->parent=parent;
        p->right=parent;
        parent->parent=p;
        p->parent=grandpa;
        if(grandpa==NULL)root=p;
        else{
            if(grandpa->left==parent)
            grandpa->left=p;
            else grandpa->right=p;
        }
        up_son(parent);up_son(p);
    }
    void remove(node* p,int key){
        if(key<p->val)remove(p->left,key);
        else if(key>p->val)remove(p->right,key);
        else {
            if(p->right==NIL){
                remove_one(p);
                return;
            }
            node *replace=suf(p);
            std::swap(p->val,replace->val);
            remove_one(replace);
        }
        up_son(p);
    }
    void remove_one(node *p){
        node *child=p->left==NIL?p->right:p->left;
        if(p->parent==NULL&&p->left==NIL&&p->right==NIL){
            p=NULL;
            root=p;return;
        }
        if(p->parent==NULL){
            delete p;
            child->parent=NULL;
            root=child;
            set_black(root);return;
        }
        if(lson(p))p->parent->left=child;
        else p->parent->right=child;
        child->parent=p->parent;
        if(is_black(p)){
            if(is_red(child))set_black(child);
            else up_remove(child);
        }
        delete p;
    }
    void up_remove(node *p){
        if(p->parent==NULL){
            set_black(p);
            return;
        }
        if(is_red(p->sibling())){
             set_red(p->parent);
             set_black(p->sibling());
             if(lson(p))rotate_left(p->parent);
             else rotate_right(p->parent);
        }
        if(is_black(p->parent)&&is_black(p->sibling()) \
        &&is_black(p->sibling()->left)&&is_black(p->sibling()->right)){
            set_red(p->sibling());
            up_remove(p->parent);
        }else if(is_red(p->parent)&&is_black(p->sibling()) \
        &&is_black(p->sibling()->left)&&is_black(p->sibling()->right)){
            set_red(p->sibling());
            set_black(p->parent);
        }
        else{
            if(is_black(p->sibling())){
                if(lson(p)&&is_red(p->sibling()->left) \
                &&is_black(p->sibling()->right)){
                    set_red(p->sibling());
                    set_black(p->sibling()->left);
                    rotate_right(p->sibling()->left);
                }else if(rson(p)&&is_black(p->sibling()->left) \
                &&is_red(p->sibling()->right)){
                    set_red(p->sibling());
                    set_black(p->sibling()->right);
                    rotate_left(p->sibling()->right);
                }
            }
            p->sibling()->color=p->parent->color;
            set_black(p->parent);
            if(lson(p)){
                set_black(p->sibling()->right);
                rotate_left(p->sibling());
            }else{
                set_black(p->sibling()->left);
                rotate_right(p->sibling());
            }
        }
    }
public:
    vector<T> DLR(){vector<T>v;DLR(root,v);return v;}
    vector<T> LDR(){vector<T>v;LDR(root,v);return v;}
    vector<T> LRD(){vector<T>v;LRD(root,v);return v;}
    vector<T> bfs(){vector<T>v;bfs(root,v);return v;}
private:
    void bfs(node *Tree,vector<T>&v)
    {
        if(Tree==NULL)return;
        queue<node*>q;
        q.push(Tree);
        while(!q.empty())
        {
            node *p=q.front();
            v.push_back(p->val);
            q.pop();
            if(p->left!=NIL)q.push(p->left);
            if(p->right!=NIL)q.push(p->right);
        }
    }
    void DLR(node *Tree,vector<T>&v)
    {
        if(Tree==NULL||Tree==NIL){return ;}
        v.push_back(Tree->val);
        DLR(Tree->left,v);
        DLR(Tree->right,v);
    }
    void LDR(node *Tree,vector<T>&v)
    {
        if(Tree==NULL||Tree==NIL){return ;}
        LDR(Tree->left,v);
        v.push_back(Tree->val);
        LDR(Tree->right,v);
    }
    void LRD(node *Tree,vector<T>&v)
    {
        if(Tree==NULL||Tree==NIL){return ;}
        LRD(Tree->left,v);
        LRD(Tree->right,v);
        v.push_back(Tree->val);
    }
private:
    node *root,*NIL;
    int n;
};