red black tree

insert-fixup

总共3种情况

我们讨论待插入结点\(x\),\(x\)的父亲\(p\),\(x\)的祖父\(g\)和\(g\)的另一儿子\(u\)(\(x\)的叔节点)

如果在若干次操作后\(p\)的颜色为黑色，则已经调整完毕，可以退出

否则判断3种情况：

case 1

\(u\)为红色

图中\(z\)是\(x\).

那么\(p\),\(u\)都为红色，\(g\)为黑色。

将三者反色，此时红色转移到了\(g\)上，需要继续递归判断\(g\)是否有双红

case 2

\(u\)黑色且\(x\)与\(p\)异侧(指一个是左孩子，另一个是右孩子，\(x\),\(p\),\(g\)不是一条直线)

此时将\(x\)旋上来，转化为情况3

case 3

\(u\)为黑色且\(x\),\(p\)同侧

那么把\(p\)旋上来,\(g\)变为\(p\)的儿子

将\(p\)变黑，\(g\)变红

此时可以做到\(p\)的所有儿子黑高相同，并且无双红，完成调整，可以退出

code:

template<typename T> void red_black_tree<T>::insert_fixup(int x){
    while(a[x].fa && a[a[x].fa].col==RED){
        int p=a[x].fa,g=a[p].fa,k=gets(x),t=gets(p),u=a[g].son[t^1];//不用判断g是否存在，因为父亲为red则父亲一定不是根，父亲一定有父亲，g一定存在
        if(a[u].col==RED){//case 1
            a[p].col=a[u].col=BLACK;
            a[g].col=RED;
            x=g;
            continue;//下一循环判断
        }
        if(k!=t) rotate(x),swap(x,p);// case 2
        // case 3
        a[p].col=BLACK,a[g].col=RED;
        rotate(p);
    }
    a[rt].col=BLACK;
}

• 注： 这里的为了不判断旋转方向，\(rotate(x)\) 的 \(x\) 是子节点而不是父亲，这个函数把 \(x\) 旋到它的父亲的位置

delete

先找到待删除的节点，然后用它的后继节点代替他。

实现时，需要判断2种情况。

一种是左/右子树为空，直接用儿子替代他

另一种是两颗子树都不空，这时要找到后继节点，把它替换上来

定义\(transplant(u,v)\)函数为用\(v\)节点替代\(u\)

\(minimum(u)\)函数为找到\(u\)的后继节点

template<typename T> void red_black_tree<T>::transplant(int u,int v){
    if(u==rt) rt=v;
    else a[a[u].fa].son[gets(u)]=v;
    a[v].fa=a[u].fa;
}
template<typename T> int red_black_tree<T>::minimum(int x){
    while(a[x].son[0]) x=a[x].son[0];
    return x;
}
template<typename T> void red_black_tree<T>::del(int z){
    int y=z,x=0;
    color col=a[y].col;
    if(!a[z].son[0]) x=a[z].son[1],transplant(z,a[z].son[1]);
    else if(!a[z].son[1]) x=a[z].son[0],transplant(z,a[z].son[0]);
    else{
        y=minimum(a[z].son[1]);
        col=a[y].col;
        x=a[y].son[1];
        a[x].fa=y;//如果x为空，这一步可以保证update(a[x].fa)的正确性
        if(a[y].fa!=z){
            transplant(y,a[y].son[1]);
            a[y].son[1]=a[z].son[1];
            a[a[y].son[1]].fa=y;
        }
        transplant(z,y);
        a[y].son[0]=a[z].son[0];
        a[a[y].son[0]].fa=y;
        a[y].col=a[z].col;
    }
    update(a[x].fa);//不要从x开始更新，x可能为空
    if(col==BLACK) delete_fixup(x);
}

delete_fixup

如果当前待调整节点\(x\)颜色为红，则停止调整 (结束后染成黑色)

如果\(x\)为根，直接把\(x\)设为黑色

否则判断4种情况：

考虑待调整节点\(x\), \(x\)的父亲\(p\), \(x\)的兄弟节点\(w\)和兄弟节点的两个儿子

我们让\(x\)指向的节点带上一层额外的黑色，并在最终调整结束时将\(x\)染黑，从而平衡删除节点带来的黑高\(-1\)。

case 1

\(w\)为红色

此时将\(w\)旋上来，将\(p\),\(w\)反色，此时其他子树的黑高不变，\(x\)黑高仍然\(-1\), 此时转化为 case 2,3,4

case 2

\(w\)和他的两个儿子都是黑色，此时将\(w\)染黑，将\(p\)（图中的\(c\)）设为新的\(x\).

观察到原本的\(x\)和\(w\)的子树黑高已经相等，如果\(p\)为红色，直接把\(p\)染黑, 即可让这个子树缺失的一个黑高补上，即可完成调整。

否则重新进入判断

• 注： 判断\(p\)为红色是下一循环要做的条件判断，如果满足就可以退出循环

case 3

\(w\)黑色，与\(w\)异侧的儿子为红色

此时将异侧儿子旋上来，\(w\)和儿子反色，进入case 4

case 4

\(w\)黑色，与\(w\)同侧的儿子为红色

此时将\(w\)旋上来，\(w\),\(w\)的同侧儿子和\(p\)反色

观察到\(x\)分支由于\(p\)变黑，从根到\(x\)所有\(NIL\)的黑高\(+1\)，从而平衡

\(w\)异侧儿子子树黑高不变

\(w\)异侧儿子由于变黑，子树黑高不变

此时调整完毕，将\(x\)设为根即可退出

code:

template<typename T> void red_black_tree<T>::delete_fixup(int x){
    while(x!=rt && a[x].col==BLACK){
        int k=gets(x),w=a[a[x].fa].son[k^1];
        if(a[w].col==RED){ //case 1
            a[w].col=BLACK;
            a[a[x].fa].col=RED;
            rotate(w);
            w=a[a[x].fa].son[k^1];
        }
        if(a[a[w].son[0]].col==BLACK && a[a[w].son[1]].col==BLACK){ //case 2
            a[w].col=RED;
            x=a[x].fa;
        }else{ 
            if(a[a[w].son[k^1]].col==BLACK){//case 3
                a[a[w].son[k]].col=BLACK;
                a[w].col=RED; 
                rotate(a[w].son[k]);
                w=a[a[x].fa].son[k^1];
            }
            //case 4
            a[w].col=a[a[x].fa].col;
            a[a[x].fa].col=BLACK;
            a[a[w].son[k^1]].col=BLACK;
            rotate(w);
            x=rt;
        }
    }
    a[x].col=BLACK;
}

证明：红黑树的删除操作中，旋转的次数是 \(O(1)\) 的

分情况

如果遇到情况 \(1\), 那么转化为情况 \(2, 3, 4\), 并且注意到情况 \(2\) 的根是红的，所以会直接结束；如果转化为情况 \(3, 4\), 也会在 \(2\) 次操作内结束，所以旋转至多 \(3\) 次；

如果一开始遇到情况 \(2\), 那么没有旋转，重新递归判断情况；

如果遇到 \(3, 4\), 那么 \(2\) 次旋转内结束

综上，旋转次数是 \(O(1)\)

完整代码

#include<bits/stdc++.h>
using namespace std;
const int N=2e5+10,INF=1e9+10;
int T,n,cnt;
template<typename T> class red_black_tree{
    public:
        int num,rt;
        void Init();
        int build(T val,int fa);
        void Insert(T val);
        void Delete(T val);
        int getrank(int x, T val);
        T getval(int x, int rank);
        int getpre(T val);
        int getnxt(T val);
        #ifdef __RED_BLACK_TREE_DEBUG__ 
            void dfs(int x){
                if(x==NIL) return;
                cout<<x<<"("<<(a[x].col==BLACK?"BLACK":"RED")<<","<<a[x].val<<","<<a[x].sz<<")->"<<a[x].son[0]<<","<<a[x].son[1]<<endl;
                dfs(a[x].son[0]),dfs(a[x].son[1]);
            }
        #endif
    private:
        const int NIL = 0;
        enum color{RED,BLACK};
        struct node{
            int fa,son[2];
            T val;
            int sz,cnt;
            color col;
        };
        node a[N];    

        void pushup(int x);
        int gets(int x);
        void rotate(int x);
        void insert_fixup(int x);
        void delete_fixup(int x);
        void update(int x);
        void transplant(int u,int v);
        int minimum(int x);
        void del(int z);
};
template<typename T> void red_black_tree<T>::Init(){
    a[NIL].col=BLACK;
}
template<typename T> void red_black_tree<T>::pushup(int x){
    a[x].sz=a[a[x].son[0]].sz+a[x].cnt+a[a[x].son[1]].sz;
}
template<typename T> int red_black_tree<T>::gets(int x){
    return a[a[x].fa].son[1]==x;
}
template<typename T> void red_black_tree<T>::rotate(int x){
    int y=a[x].fa,z=a[y].fa,k=gets(x),t=gets(y);
    a[y].son[k]=a[x].son[k^1];
    if(a[x].son[k^1]!=NIL) a[a[x].son[k^1]].fa=y;
    a[x].son[k^1]=y;a[y].fa=x;
    if(z!=NIL) a[z].son[t]=x;
    a[x].fa=z;
    pushup(y),pushup(x);
    if(rt==y) rt=x;//换根
}
template<typename T> int red_black_tree<T>::build(T val,int fa){
    a[++num].val=val;
    a[num].col=RED;
    a[num].cnt=a[num].sz=1;
    a[num].fa=fa;
    a[num].son[0]=a[num].son[1]=NIL;
    return num;
}
template<typename T> void red_black_tree<T>::insert_fixup(int x){
    while(a[x].fa && a[a[x].fa].col==RED){
        int p=a[x].fa,g=a[p].fa,k=gets(x),t=gets(p),u=a[g].son[t^1];//不用判断g是否存在，因为父亲为red则父亲一定不是根，父亲一定有父亲，g一定存在
        if(a[u].col==RED){//case 1
            a[p].col=a[u].col=BLACK;
            a[g].col=RED;
            x=g;
            continue;//下一循环判断
        }
        if(k!=t) rotate(x),swap(x,p);// case 2
        // case 3
        a[p].col=BLACK,a[g].col=RED;
        rotate(p);
    }
    a[rt].col=BLACK;
}
template<typename T> void red_black_tree<T>::delete_fixup(int x){
    while(x!=rt && a[x].col==BLACK){
        int k=gets(x),w=a[a[x].fa].son[k^1];
        if(a[w].col==RED){ //case 1
            a[w].col=BLACK;
            a[a[x].fa].col=RED;
            rotate(w);
            w=a[a[x].fa].son[k^1];
        }
        if(a[a[w].son[0]].col==BLACK && a[a[w].son[1]].col==BLACK){ //case 2
            a[w].col=RED;
            x=a[x].fa;
        }else{ 
            if(a[a[w].son[k^1]].col==BLACK){//case 3
                a[a[w].son[k]].col=BLACK;
                a[w].col=RED; 
                rotate(a[w].son[k]);
                w=a[a[x].fa].son[k^1];
            }
            //case 4
            a[w].col=a[a[x].fa].col;
            a[a[x].fa].col=BLACK;
            a[a[w].son[k^1]].col=BLACK;
            rotate(w);
            x=rt;
        }
    }
    a[x].col=BLACK;
}

template<typename T> void red_black_tree<T>::update(int x){
    while(x!=NIL) pushup(x),x=a[x].fa;
}
template<typename T> void red_black_tree<T>::Insert(T val){
    int now=rt,fa=0;
    while(now && val!=a[now].val) fa=now,now=a[now].son[val>a[now].val];
    if(now) a[now].cnt++;
    else{
        now=build(val,fa);
        if(fa) a[fa].son[val>a[fa].val]=now;
        else rt=now;
    }
    update(now);
    insert_fixup(now);
}

template<typename T> void red_black_tree<T>::transplant(int u,int v){
    if(u==rt) rt=v;
    else a[a[u].fa].son[gets(u)]=v;
    a[v].fa=a[u].fa;
}
template<typename T> int red_black_tree<T>::minimum(int x){
    while(a[x].son[0]) x=a[x].son[0];
    return x;
}
template<typename T> void red_black_tree<T>::del(int z){
    int y=z,x=0;
    color col=a[y].col;
    if(!a[z].son[0]) x=a[z].son[1],transplant(z,a[z].son[1]);
    else if(!a[z].son[1]) x=a[z].son[0],transplant(z,a[z].son[0]);
    else{
        y=minimum(a[z].son[1]);
        col=a[y].col;
        x=a[y].son[1];
        a[x].fa=y;
        if(a[y].fa!=z){
            transplant(y,a[y].son[1]);
            a[y].son[1]=a[z].son[1];
            a[a[y].son[1]].fa=y;
        }
        transplant(z,y);
        a[y].son[0]=a[z].son[0];
        a[a[y].son[0]].fa=y;
        a[y].col=a[z].col;
    }
    update(a[x].fa);
    if(col==BLACK) delete_fixup(x);
}
template<typename T> void red_black_tree<T>::Delete(T val){
    int now=rt;
    while(now && val!=a[now].val) now=a[now].son[val>a[now].val];
    if(now){
        if(a[now].cnt>1) a[now].cnt--,update(now);
        else del(now);
    }
}


template<typename T> int red_black_tree<T>::getrank(int x, T val){
    if(x==NIL) return 1;
    if(val<a[x].val) return getrank(a[x].son[0], val);
    else if(val>a[x].val) return a[x].cnt+a[a[x].son[0]].sz+getrank(a[x].son[1], val);
    else{
        return a[a[x].son[0]].sz+1;
    }
}
template<typename T> T red_black_tree<T>::getval(int x, int rank){
    if(x==NIL) return INF;
    if(rank<=a[a[x].son[0]].sz) return getval(a[x].son[0], rank);
    else if(rank>a[a[x].son[0]].sz+a[x].cnt) return getval(a[x].son[1],rank-(a[a[x].son[0]].sz+a[x].cnt));
    else return a[x].val;
}
template<typename T> int red_black_tree<T>::getpre(T val){
    int now=rt,pre=0;
    while(now!=NIL){
        if(val<=a[now].val) now=a[now].son[0];
        else if(val>a[now].val) pre=now,now=a[now].son[1];
    }
    return a[pre].val;
}
template<typename T> int red_black_tree<T>::getnxt(T val){
    int now=rt,nxt=0;
    while(now!=NIL){
        if(val<a[now].val) nxt=now,now=a[now].son[0];
        else if(val>=a[now].val) now=a[now].son[1];
    }
    return a[nxt].val;
}

red_black_tree<long long> S;
void solve() {
    cin>>n;
    S.Init();
    S.Insert(INF),S.Insert(-INF);
    for(int i=1,op,x;i<=n;++i){
        cin>>op>>x;
        if(op==1){
            S.Insert(x);
        }else if(op==2){
            S.Delete(x);
        }else if(op==3){
            cout<<S.getrank(S.rt, x)-1<<endl;
        }else if(op==4){
            cout<<S.getval(S.rt, x+1)<<endl;
        }else if(op==5){
            cout<<S.getpre(x)<<endl;
        }else if(op==6){
            cout<<S.getnxt(x)<<endl;   
        }
    }
}
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr),cout.tie(nullptr);
    solve();
    return 0;
}

等等，上面这份代码是错的

问题出在了\(NIL\)上，由于本人不喜欢用指针，所以当程序访问到\(NIL\)时，他不会报错，而是会继续运行

也就是说，像\(pushup\)这样的函数在访问到\(NIL\)节点时，会将\(NIL\)的\(size\)算上，而在上面这份代码中，\(NIL\)的大小会随着程序运行变成非零

血的教训：写大型数据结构，还是乖乖用指针

经过查错/对拍后，下面代码的RBT与Treap,scapegoat,AVL答案一致，且成为最快（P3369中RBT 195ms, Treap 211ms ,AVL 307ms,scapegoat 346ms）

#include<bits/stdc++.h>
using namespace std;
const int N=3e6+10,INF=1e9+10;
int T,n,cnt;
#define __RED_BLACK_TREE_DEBUG__
template<typename T> class red_black_tree{
    public:
        int num,rt;
        void Init();
        int build(T val,int fa);
        void Insert(T val);
        void Delete(T val);
        int getrank(int x, T val);
        T getval(int x, int rank);
        int getpre(T val);
        int getnxt(T val);
        #ifdef __RED_BLACK_TREE_DEBUG__ 
            void dfs(int x){
                if(x==NIL) return;
                cout<<x<<"("<<(a[x].col==BLACK?"BLACK":"RED")<<","<<a[x].val<<","<<a[x].sz<<")->"<<a[x].son[0]<<","<<a[x].son[1]<<endl;
                dfs(a[x].son[0]),dfs(a[x].son[1]);
            }
        #endif
    private:
        const int NIL = 0;
        enum color{RED,BLACK};
        struct node{
            int fa,son[2];
            T val;
            int sz,cnt;
            color col;
        };
        node a[N];    

        void pushup(int x);
        int gets(int x);
        void rotate(int x);
        void insert_fixup(int x);
        void delete_fixup(int x);
        void update(int x);
        void transplant(int u,int v);
        int minimum(int x);
        void del(int z);
};
template<typename T> void red_black_tree<T>::Init(){
    a[NIL].col=BLACK;
    rt=num=0;
}
template<typename T> void red_black_tree<T>::pushup(int x){
    a[x].sz=(a[x].son[0]!=NIL)*a[a[x].son[0]].sz+a[x].cnt+(a[x].son[1]!=NIL)*a[a[x].son[1]].sz;
    //可能访问到NIL
}
template<typename T> int red_black_tree<T>::gets(int x){
    return a[a[x].fa].son[1]==x;
}
template<typename T> void red_black_tree<T>::rotate(int x){
    int y=a[x].fa,z=a[y].fa,k=gets(x),t=gets(y);
    a[y].son[k]=a[x].son[k^1];
    if(a[x].son[k^1]!=NIL) a[a[x].son[k^1]].fa=y;
    a[x].son[k^1]=y;a[y].fa=x;
    if(z!=NIL) a[z].son[t]=x;
    a[x].fa=z;
    pushup(y),pushup(x);
    if(rt==y) rt=x;//换根
}
template<typename T> int red_black_tree<T>::build(T val,int fa){
    a[++num].val=val;
    a[num].col=RED;
    a[num].cnt=a[num].sz=1;
    a[num].fa=fa;
    a[num].son[0]=a[num].son[1]=NIL;
    return num;
}
template<typename T> void red_black_tree<T>::insert_fixup(int x){
    while(a[x].fa && a[a[x].fa].col==RED){
        int p=a[x].fa,g=a[p].fa,k=gets(x),t=gets(p),u=a[g].son[t^1];//不用判断g是否存在，因为父亲为red则父亲一定不是根，父亲一定有父亲，g一定存在
        if(a[u].col==RED){//case 1
            a[p].col=a[u].col=BLACK;
            a[g].col=RED;
            x=g;
            continue;//下一循环判断
        }
        if(k!=t) rotate(x),swap(x,p);// case 2
        // case 3
        a[p].col=BLACK,a[g].col=RED;
        rotate(p);
    }
    a[rt].col=BLACK;
}
template<typename T> void red_black_tree<T>::delete_fixup(int x){
    while(x!=rt && a[x].col==BLACK){
        int k=gets(x),w=a[a[x].fa].son[k^1];
        if(a[w].col==RED){ //case 1
            a[w].col=BLACK;
            a[a[x].fa].col=RED;
            rotate(w);
            w=a[a[x].fa].son[k^1];
        }
        if(a[a[w].son[0]].col==BLACK && a[a[w].son[1]].col==BLACK){ //case 2
            a[w].col=RED;
            x=a[x].fa;
        }else{ 
            if(a[a[w].son[k^1]].col==BLACK){//case 3
                a[a[w].son[k]].col=BLACK;
                a[w].col=RED; 
                rotate(a[w].son[k]);
                w=a[a[x].fa].son[k^1];
            }
            //case 4
            a[w].col=a[a[x].fa].col;
            a[a[x].fa].col=BLACK;
            a[a[w].son[k^1]].col=BLACK;
            rotate(w);
            x=rt;
        }
    }
    a[x].col=BLACK;
}

template<typename T> void red_black_tree<T>::update(int x){
    while(x!=NIL) pushup(x),x=a[x].fa;//实际好多人这么写，我看有的库里的RBT也是这么写的，那就这样吧
}
template<typename T> void red_black_tree<T>::Insert(T val){
    int now=rt,fa=0;
    while(now && val!=a[now].val) fa=now,now=a[now].son[val>a[now].val];
    if(now) a[now].cnt++,update(now);
    else{
        now=build(val,fa);
        if(fa) a[fa].son[val>a[fa].val]=now;
        else rt=now;
        update(now);
        insert_fixup(now);//把这两句挪进来，因为只改变cnt不需要双红修正，这个逻辑谬误找了0.5h
    }
}

template<typename T> void red_black_tree<T>::transplant(int u,int v){
    if(u==rt) rt=v;
    else a[a[u].fa].son[gets(u)]=v;
    a[v].fa=a[u].fa;
}
template<typename T> int red_black_tree<T>::minimum(int x){
    while(a[x].son[0]) x=a[x].son[0];
    return x;
}
template<typename T> void red_black_tree<T>::del(int z){
    int y=z,x=0;
    color col=a[y].col;
    if(!a[z].son[0]) x=a[z].son[1],transplant(z,a[z].son[1]);
    else if(!a[z].son[1]) x=a[z].son[0],transplant(z,a[z].son[0]);
    else{
        y=minimum(a[z].son[1]);
        col=a[y].col;
        x=a[y].son[1];
        a[x].fa=y;
        if(a[y].fa!=z){
            transplant(y,a[y].son[1]);
            a[y].son[1]=a[z].son[1];
            a[a[y].son[1]].fa=y;
        }
        transplant(z,y);
        a[y].son[0]=a[z].son[0];
        a[a[y].son[0]].fa=y;
        a[y].col=a[z].col;
    }
    update(a[x].fa);
    if(col==BLACK) delete_fixup(x);
}
template<typename T> void red_black_tree<T>::Delete(T val){
    int now=rt;
    while(now && val!=a[now].val) now=a[now].son[val>a[now].val];
    if(now){
        if(a[now].cnt>1) a[now].cnt--,update(now);
        else del(now);
    }
}
template<typename T> int red_black_tree<T>::getrank(int x, T val){
    if(x==NIL) return 1;
    int szl=a[x].son[0]?a[a[x].son[0]].sz:0;//会访问到NIL
    if(val<a[x].val) return getrank(a[x].son[0], val);
    else if(val>a[x].val) return a[x].cnt+szl+getrank(a[x].son[1], val);
    else{
        return szl+1;
    }
}
template<typename T> T red_black_tree<T>::getval(int x, int rank){
    if(x==NIL) return INF;
    int szl=a[x].son[0]?a[a[x].son[0]].sz:0;//会访问到NIL
    if(rank<=szl) return getval(a[x].son[0], rank);
    else if(rank>szl+a[x].cnt) return getval(a[x].son[1],rank-(szl+a[x].cnt));
    else return a[x].val;
}
template<typename T> int red_black_tree<T>::getpre(T val){
    int now=rt,pre=0;
    while(now!=NIL){
        if(val<=a[now].val) now=a[now].son[0];
        else if(val>a[now].val) pre=now,now=a[now].son[1];
    }
    return a[pre].val;
}
template<typename T> int red_black_tree<T>::getnxt(T val){
    int now=rt,nxt=0;
    while(now!=NIL){
        if(val<a[now].val) nxt=now,now=a[now].son[0];
        else if(val>=a[now].val) now=a[now].son[1];
    }
    return a[nxt].val;
}

red_black_tree<int> S;
void solve() {
    scanf("%d",&n);
    S.Init();
    S.Insert(INF),S.Insert(-INF);
    for(int i=1,op,x;i<=n;++i){
        scanf("%d%d",&op,&x);
        if(op==1){
            S.Insert(x);
        }else if(op==2){
            S.Delete(x);
        }else if(op==3){
           printf("%d\n",S.getrank(S.rt, x)-1);
        }else if(op==4){
            printf("%d\n",S.getval(S.rt, x+1));
        }else if(op==5){
            printf("%d\n",S.getpre(x));
        }else if(op==6){
            printf("%d\n",S.getnxt(x));   
        }
    }
}
int main() {
    solve();
    return 0;
}

• 参考/图源：算法导论