【bzoj3224】平衡树模板（Splay）

bzoj 3224 普通平衡树 包含此模板全部操作，可以提交至此。

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
#include <vector>
#include <queue>
#include <cmath>
using namespace std;
const int maxn=500005;
int fa[maxn],ch[maxn][2],dat[maxn],cnt[maxn],siz[maxn],sz,root;
//fa[i]表示i的父结点，ch[i][0]表示i的左儿子，ch[i][1]表示i的右儿子，
//dat[i]表示i的关键字（即结点i代表的那个数字），cnt[i]表示i结点的关键字出现的次数（相当于权值），
//siz[i]表示包括i的这个子树的大小；sz为整棵树的大小，root为整棵树的根
int n;
inline void clear(int x){//将当前点的各项值都清0
    ch[x][0]=ch[x][1]=fa[x]=cnt[x]=dat[x]=siz[x]=0;
}
inline int get(int x){//判断当前点是它父结点的左儿子还是右儿子
    return ch[fa[x]][1]==x;
    //左儿子返回0，右儿子返回1;
}
inline void update(int x){//更新当前点的siz值 
    if(x){
        siz[x]=cnt[x];//定为该点的权值
        if(ch[x][0])siz[x]+=siz[ch[x][0]];//左儿子的大小
        if(ch[x][1])siz[x]+=siz[ch[x][1]];//右儿子的大小
    }
}
inline void rotate(int x){//将x结点rotate到它的父亲的位置。
    int old=fa[x],oldf=fa[old],wh=get(x);
//找出x与父亲的位置关系
    ch[old][wh]=ch[x][wh^1];
    fa[ch[old][wh]]=old;
    fa[old]=x;
    ch[x][wh^1]=old;
//将原先父亲更新到x的对应儿子的位置，将x原先对应儿子更新到x原先父亲的对应儿子;
    fa[x]=oldf;
    if(oldf){
        ch[oldf][ch[oldf][1]==old]=x;
    }
//将x更新到原先父亲的位置;
    update(old);
    update(x);
//更新子树大小有变化的节点。
} splay(int x){//将x结点rotate到根。
    for(int f;f=fa[x];rotate(x)){
        if(fa[f])rotate((get(x)==get(f)?f:x));
        //分类讨论，如果x与父亲，祖父三点共线，则需要rotate父亲，否则rotate x，避免出现平衡树失衡，可以尝试手动证明。
    }
    root=x;//更新根节点
}
inline void insert(int v){//插入值为v的点
    if(root==0){
//子树为空则特殊处理
        sz++;ch[sz][0]=ch[sz][1]=fa[sz]=0;
        dat[sz]=v;cnt[sz]=1;siz[sz]=1;root=sz;
        return;
    }
    int now=root,f=0;
    while(1){
        if(dat[now]==v){
            //若存在v，则增加该节点权值
            cnt[now]++;
            update(now);
            update(f);
            splay(now);
            break;
        }
        f=now;
        now=ch[now][dat[now]<v];
        if(now==0){
            //若不存在，则新增节点。
            sz++;
            ch[sz][0]=ch[sz][1]=0;
            dat[sz]=v;
            siz[sz]=1;
            cnt[sz]=1;
            fa[sz]=f;
            ch[f][dat[f]<v]=sz;
            update(f);
            splay(sz);
            break;
        }
    }
}e int find(int v){//查询v的排名
    int ans=0,now=root;
    while(1){
        if(v<dat[now])now=ch[now][0];
        else{
            ans+=(ch[now][0]?siz[ch[now][0]]:0);
            if(v==dat[now]){splay(now);return ans+1;}
            ans+=cnt[now];
            now=ch[now][1];
        }
    }
}
inline int findx(int x){//找到排名为x的点
    int now=root;
    while(1){
        if(ch[now][0]&&x<=siz[ch[now][0]]){
            now=ch[now][0];
        }else{
            int tmp=(ch[now][0]?siz[ch[now][0]]:0)+cnt[now];
            if(x<=tmp)return dat[now];
            x-=tmp;
            now=ch[now][1];
        }
    }
}
inline int pre(){//查询树上的前驱
    int now=ch[root][0];
    while(ch[now][1])now=ch[now][1];
    return now;
}
inline int nxt(){//查询树上的后继
    int now=ch[root][1];
    while(ch[now][0])now=ch[now][0];
    return now;
}
inline void del(int x){//删除节点x
    int wh=find(x);//把x旋转到根
    //分类讨论每一种情况
    if(cnt[root]>1){cnt[root]--;return;}
    if(!ch[root][0]&&!ch[root][1]){clear(root);root=0;return;}
    if(!ch[root][0]){
        int oldr=root;
        root=ch[root][1];
        fa[root]=0;
        clear(oldr);
        return;
    }else if(!ch[root][1]){
        int oldr=root;
        root=ch[root][0];
        fa[root]=0;
        clear(oldr);
        return;
    }
    //若两个儿子都有，则把x的前驱splay到树根，再把原来根节点的右儿子接到现在的根上，此时原先的根就变成了叶子节点，直接删除即可。
    int lb=pre();
    int oldr=root;
    splay(lb);
    fa[ch[oldr][1]]=root;
    ch[root][1]=ch[oldr][1];
    clear(oldr);
    update(root);
    return;
}
inline int prex(int x){//查x的前驱
    //查询x的前驱即插入x，查询树上的前驱（此时x为root），然后删除x的过程，后继同理。
    insert(x);
    int ans=pre();
    del(x);
    return dat[ans];
}
inline int nxtx(int x){//查x的后继
    insert(x);
    int ans=nxt();
    del(x);
    return dat[ans];
}
int main(){
    scanf("%d",&n);
    int opt,xx;
    while(n--){
        scanf("%d%d",&opt,&xx);
        if(opt==1){
            insert(xx);
            continue;
        }
        if(opt==2){
            del(xx);
            continue;
        }
        if(opt==3){
            printf("%d\n",find(xx));
            continue;
        }
        if(opt==4){
            printf("%d\n",findx(xx));
            continue;
        }
        if(opt==5){
            printf("%d\n",prex(xx));
            continue;
        }
        if(opt==6){
            printf("%d\n",nxtx(xx));
            continue;
        }
    }
    return 0;
}