Treap模板

Posted 2020-11-02 Zinn

题目：https://www.luogu.org/problemnew/show/P3369

Treap模板。

代码如下：

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
int const MAXN=100005,inf=1e9;
int n,rt,cnt;
struct N{
    int ch[3],val,siz,pri,num;
}t[MAXN];
int newnode()
{
    int r=++cnt;
    t[r].siz=1;t[r].num=1;t[r].pri=rand();
    return r;
}
void pushup(int &o)
{
    t[o].siz=t[t[o].ch[0]].siz+t[t[o].ch[1]].siz+t[o].num;
}
void rotate(int &o,int d)
{
    int u=t[o].ch[d];
    t[o].ch[d]=t[u].ch[d^1];
    t[u].ch[d^1]=o;
    t[u].siz=t[o].siz;
    pushup(o);
    o=u;
}
void add(int &o,int x)
{
    if(!o)
    {
        o=newnode();
        t[o].val=x;
        return;
    }
    if(t[o].val==x)
    {
        t[o].num++;t[o].siz++;
        return;
    }
    int d=t[o].val<x;
    add(t[o].ch[d],x);
    pushup(o);
    if(t[t[o].ch[d]].pri<t[o].pri)rotate(o,d);
}
void del(int &o,int x)
{
    if(!o)return;
    if(t[o].val==x)
    {
        if(t[o].num>1)t[o].num--;
        else if(t[o].ch[1]*t[o].ch[0]==0)//
            o=t[o].ch[1]+t[o].ch[0];
        else
        {
            int d=t[t[o].ch[0]].pri<t[t[o].ch[1]].pri;
            rotate(o,d^1);
            del(o,x);//o!
        }
        pushup(o);
        return;
    }
    int d=t[o].val<x;
    del(t[o].ch[d],x);
    pushup(o);
}
int query(int &o,int x)
{
    if(!o)return 0;
    if(t[o].val==x)return t[t[o].ch[0]].siz+1;
    int d=t[o].val<x;
    if(!d)return query(t[o].ch[d],x);
    return query(t[o].ch[d],x)+t[o].num+t[t[o].ch[0]].siz;
}
int qr(int &o,int x)
{
    if(!o||!x)return 0;
    int d=x-t[t[o].ch[0]].siz;
    if(d<=0)return qr(t[o].ch[0],x);
    if(d<=t[o].num)return t[o].val;
    return qr(t[o].ch[1],x-t[t[o].ch[0]].siz-t[o].num);
}
int pre(int &o,int x)
{
    if(!o)return -inf;//
    if(t[o].val<x)return max(t[o].val,pre(t[o].ch[1],x));
    return pre(t[o].ch[0],x);
}
int lst(int &o,int x)
{
    if(!o)return inf;//
    if(t[o].val>x)return min(t[o].val,lst(t[o].ch[0],x));
    return lst(t[o].ch[1],x);
}
int main()
{
    scanf("%d",&n);
    for(int i=1,x,opt;i<=n;i++)
    {
        scanf("%d%d",&opt,&x);
        if(opt==1)add(rt,x);//插入 x 数
        if(opt==2)del(rt,x);//删除 x 数(若有多个相同的数，只删除一个)
        if(opt==3)printf("%d\n",query(rt,x));
        //查询 x 数的排名(排名定义为比当前数小的数的个数 +1,若有多个相同的数,输出最小的排名)
        if(opt==4)printf("%d\n",qr(rt,x));//查询排名为 x 的数
        if(opt==5)printf("%d\n",pre(rt,x));//求 x 的前驱(前驱定义为小于 x ，且最大的数)
        if(opt==6)printf("%d\n",lst(rt,x));//求 x 的后继(后继定义为大于 x ，且最小的数)
    }
    return 0;
}