BZOJ3224: 普通平衡树

#include<bits/stdc++.h>
using namespace std;
const int inf=0x7fffffff;
const int bow=101000;
int t,root,tot;
int dat[bow],val[bow],l[bow],r[bow],size[bow],cnt[bow];
void update(int x){size[x]=size[l[x]]+size[r[x]]+cnt[x];}
int cte(int c){val[++tot]=c;dat[tot]=rand();cnt[tot]=size[tot]=1;return tot;}
void build(){cte(-inf);cte(inf);root=1;r[1]=2;update(root);}
void zig(int &p){int q=l[p];l[p]=r[q];r[q]=p;p=q;update(r[p]);update(p);}
void zag(int &p){int q=r[p];r[p]=l[q];l[q]=p;p=q;update(l[p]);update(p);}
void insert(int &p,int c)
{
    if(p==0){p=cte(c);return;}
    if(val[p]==c){cnt[p]++;update(p);return;}
    if(c<val[p]){insert(l[p],c);if(dat[p]<dat[l[p]])zig(p);}
    else {insert(r[p],c);if(dat[p]<dat[r[p]])zag(p);}
    update(p);
}
void remove(int &p,int c)
{
    if(p==0)return;
    if(val[p]==c){
    if(cnt[p]>1){cnt[p]--;update(p);return;}
    if(l[p]||r[p])
    {if(r[p]==0||dat[r[p]]<dat[l[p]])zig(p),remove(r[p],c);
    else zag(p),remove(l[p],c);
    update(p);}
    else p=0;
    return;}
    c<val[p]?remove(l[p],c):remove(r[p],c);
    update(p);
}
int grbv(int p,int c)
{
    if(p==0)return 0;
    if(c==val[p])return size[l[p]]+1;
    if(c<val[p])return grbv(l[p],c);
    return grbv(r[p],c)+cnt[p]+size[l[p]];
}
int gvbr(int p,int rk){
    if(p==0)return inf;
    if(size[l[p]]>=rk)return gvbr(l[p],rk);
    if(size[l[p]]+cnt[p]>=rk)return val[p];
    return gvbr(r[p],rk-cnt[p]-size[l[p]]);
}
int getpre(int c){
    int ans=1,p=root;
    while(p){
        if(c==val[p]){if(l[p]>0){p=l[p];while(r[p]>0)p=r[p];ans=p;}break;}
        if(val[p]<c&&val[ans]<val[p])ans=p;
        p=c<val[p]?l[p]:r[p];
    }
    return val[ans];
}
int getnxt(int c){
    int ans=2,p=root;
    while(p){
        if(c==val[p]){if(r[p]>0){p=r[p];while(l[p]>0)p=l[p];ans=p;}break;}
        if(val[p]>c&&val[ans]>val[p])ans=p;
        p=c<val[p]?l[p]:r[p];
    }
    return val[ans];
}
int main()
{
    build();
   cin>>t;
   while(t--)
   {
       int opt,c;
       cin>>opt>>c;
       if(opt==1)insert(root,c);
    if(opt==2)remove(root,c);
    if(opt==3)cout<<(grbv(root,c)-1)<<endl;
    if(opt==4)cout<<gvbr(root,c+1)<<endl;
    if(opt==5)cout<<getpre(c)<<endl;
    if(opt==6)cout<<getnxt(c)<<endl;
    }    
}