P3369 【模板】普通平衡树------平衡树Treap || FHQ_treap

原题
一道典型的平衡树，
cnt[]是这个点出现的次数，siz[]是这个子树中有节点的数量(包括一样的节点)，key[]记录该节点的数值，fa[]记录该节点的父节点

#include<bits/stdc++.h>
using namespace std;
const int MX=1e5+9;
int n,tot=0,root,x,order;
int fa[MX],son[MX][2],key[MX],cnt[MX],siz[MX];

int get(int x){
    return son[fa[x]][1]==x;   // 判断x是不是右子树，是返回1，左子树返回0
}

int pushup(int x){
    siz[x]=siz[son[x][0]]+siz[son[x][1]]+cnt[x];  // 一定要把这个节点的cnt[]加上
}

void rotate(int x){  // 该旋转是将x与x的父树换位置(深度-1)，并且将x的其中一个子树插到原来x再f的位置上
    int f=fa[x],ff=fa[f],w=get(x);  
    ff?son[ff][get(f)]=x:root=x;fa[x]=ff;   // 如果ff=0，那么必须要给root赋值x
    son[f][w]=son[x][!w];fa[son[x][!w]]=f;  // f的子树x的位置，现在要插一个x的子树
    son[x][!w]=f,fa[f]=x;  // f现在是x的子树，
    pushup(f);  // 先更新f，再更新x，因为经过旋转之后，x是f的子树，而且，这样时间会减少一些
    pushup(x);
}

void sply(int x,int i){  // 将x旋转到i的左右某个子树中，就相当于x在i的子树中是置顶了
    while( fa[x]!=i ){
        int f=fa[x],ff=fa[f];
        if( ff==i )  // 如果ff已经是i的话，就只需要再旋转一次了
            rotate(x);
        else{
            if( get(f)==get(x) )
                rotate(f),rotate(x);   // 三点一线情况下，要先转f，再转x
            else
                rotate(x),rotate(x);  // 转2次x，
        }
    }
    return ;
}

void inst(int &rt,int x){   //重点，这里要引用，因为会改变数值
    if( !rt ){
        rt=++tot;
        key[rt]=x,siz[rt]=cnt[rt]=1;
        return ;
    }
    if( key[rt]==x ){
        siz[rt]++,cnt[rt]++;
        return ;
    }
    if( key[rt]>x )
        inst(son[rt][0],x),fa[son[rt][0]]=rt;
    if( key[rt]<x )
        inst(son[rt][1],x),fa[son[rt][1]]=rt;
    pushup(rt);
}

int que_trank(int rt,int x){
    while( 1 ){
        if( key[rt]==x ){
            sply(rt,0);   // 这样使用简单明了
            return siz[son[rt][0]]+1;
        }
        if( key[rt]<x )
            rt=son[rt][1];
        else
            rt=son[rt][0];
    }
}

int que_order(int rt,int x){
    while( 1 ){
        if( siz[son[rt][0]]+1<=x && x<=siz[son[rt][0]]+cnt[rt] )
            return key[rt];
        if( x<siz[son[rt][0]]+1 )
            rt=son[rt][0];
        else{
            x-=(siz[son[rt][0]]+cnt[rt]);   // 这2行不可以换位置
            rt=son[rt][1];
        }
    }
}

int quepre(int rt,int x){   //  先找到第一个比x小的，再一直往右子树走即可
    int ans=0;
    while( rt ){
        if( key[rt]>=x  )   
            rt=son[rt][0];
        else{
            ans=rt;
            rt=son[rt][1];
        }
    }
    return ans;
}

int quelast(int rt,int x){  //同上
    int ans=0;
    while( rt ){
        if( key[rt]<=x )
            rt=son[rt][1];
        else{
            ans=rt;
            rt=son[rt][0];
        }
    }
    return ans;
}

int que_min(int rt){
    int ans=-1;
    while( rt ){
        ans=rt;
        rt=son[rt][0];
    }
    return ans;
}

void del(int rt,int x){
    if( key[rt]==x ){
        if( cnt[rt]>1 )
            cnt[rt]--,siz[rt]--;  // 如果已经存在不止一个，就直接-1即可
        else{
            sply(rt,0);   // 把rt放到树顶
            int p=que_min(son[rt][1]);
            if( ~p ){    // 如果key[rt]不是最大的
                sply(p,rt);  // 把p放在rt的右子树中，
                fa[p]=0,root=p;  // 然后把rt的左子树放在p的左子树中(就是把rt这个节点个给顶替掉)
                son[p][0]=son[rt][0],fa[son[rt][0]]=p;  
                pushup(p);   // 更新p节点的大小，因为p插了一个rt的左节点
            }
            else   // 如果key[rt]是最大的，那么将左子树顶替rt的位置即可
                fa[son[rt][0]]=0,root=son[rt][0];   // 记得初始化
        }
        return ;   // 这个return不可以忘记
    }
    if( key[rt]<x )
        del(son[rt][1],x);
    else
        del(son[rt][0],x);
    pushup(rt);
}

int main(){
 //   freopen("input.txt","r",stdin);
    scanf("%d",&n);
    while( n-- ){
        scanf("%d %d",&order,&x);
        if( order==1 ) inst(root,x),sply(tot,0),root=tot;  // root不能忘记，因为只插入一个点时，root不会被赋值
        else if( order==2 ) del(root,x);
        else if( order==3 ) printf("%d\n",que_trank(root,x));
        else if( order==4 ) printf("%d\n",que_order(root,x));
        else if( order==5 ) printf("%d\n",key[quepre(root,x)]);
        else printf("%d\n",key[quelast(root,x)]);
    }
    return 0;
}

也可以使用fhq_treap这种优化的数据结构写，并且此时的代码更加简洁：

#include<bits/stdc++.h>
using namespace std;
const int MX=1e5+9;
int tot=0,Root;
int val[MX],siz[MX],son[MX][2],ran[MX];

void pushup(int k){
    siz[k]=siz[son[k][1]]+siz[son[k][0]]+1;
}

void split(int o,int k,int &x,int &y){
    if( !o )
        x=y=0;
    else{
        if( val[o]<=k )
            x=o,split(son[o][1],k,son[o][1],y);
        else
            y=o,split(son[o][0],k,x,son[o][0]);
        pushup(o);
    }
}

int merge(int x,int y){
    if( !x || !y )
        return x+y;
    else{
        if( ran[x]>ran[y] ){
            son[x][1]=merge(son[x][1],y);
            pushup(x);
            return x;
        }
        else{
            son[y][0]=merge(x,son[y][0]);
            pushup(y);
            return y;
        }
    }
}

int newnode(int k){
    siz[++tot]=1;
    val[tot]=k;
    ran[tot]=rand();
    return tot;
}

void insert(int k){
    int x,y;
    split(Root,k,x,y);
    Root=merge(merge(x,newnode(k)),y);
}

void remove(int k){
    int x,y,z;
    split(Root,k,x,y);
    split(x,k-1,x,z);
    z=merge(son[z][0],son[z][1]);
    Root=merge(merge(x,z),y);
}

void kthrank(int k){
    int x,y;
    split(Root,k-1,x,y);
    printf("%d\n",siz[x]+1);
    merge(x,y);
}

int rank(int o,int k){
    if( siz[son[o][0]]>=k )
        return rank(son[o][0],k);
    else if( siz[son[o][0]]+1==k )
        return val[o];
    else
        return rank(son[o][1],k-siz[son[o][0]]-1);
}

void pre(int k){
    int x,y;
    split(Root,k-1,x,y);
    printf("%d\n",rank(x,siz[x]));
    merge(x,y);
}

void last(int k){
    int x,y;
    split(Root,k,x,y);
    printf("%d\n",rank(y,1));
    merge(x,y);
}

int main()
{
    //freopen("input.txt","r",stdin);
    int n;
    scanf("%d",&n);
    while( n-- ){
        int order,x;
        scanf("%d %d",&order,&x);
        if( order==1 ) insert(x);
        else if( order==2 ) remove(x);
        else if( order==3 ) kthrank(x);
        else if( order==4 ) printf("%d\n",rank(Root,x));
        else if( order==5 ) pre(x);
        else last(x);
    }
    return 0;
}