Treap(四)——#104. 普通平衡树

题目链接：https://loj.ac/problem/104
解题思路
Treap模板题，涉及插入，删除，找前驱，后继，找第k小，找某数排名。
AC代码

//Treap的一些基本操作
#include <iostream>
#include <stdio.h>
#include <algorithm>
using namespace std;
const int inf=0x3f3f3f3f;
const int N=5e5+5;
int rt,tot;
struct node
{
    int lc,rc,key,pri,cnt,sze;
    #define lc(x) t[x].lc //左子结点,重定义以后,下文lc(x)就是t[x].lc
    #define rc(x) t[x].rc //右子结点
    #define v(x) t[x].key //权值
    #define p(x) t[x].pri //优先级
    #define c(x) t[x].cnt //有重复的元素,在结点上再加一个域,记录重复结点的个数
    #define s(x) t[x].sze //以x为根的子树结点个数,即size
}t[N];
inline void upt(const int &k)
{
    s(k)=s(lc(k))+s(rc(k))+c(k);
}
inline void Zig(int &k)//右旋,这里加了个引用&,目的是更加方便地处理父子关系
{
    int y=lc(k);//选取当前结点k的左子结点
    lc(k)=rc(y);
    rc(y)=k;//交换y,k的父子位置,并维持BST性质
    s(y)=s(k);//维护子树的size
    upt(k);//更新结点k的size
    k=y;
}
inline void Zag(int &k)//左旋
{
    int y=rc(k);//选取当前结点k的右子结点
    rc(k)=lc(y);
    lc(y)=k;//交换y,k的父子位置,并维持BST性质
    s(y)=s(k);//维护子树的size
    upt(k);//更新结点k的size
    k=y;
}
inline void insert(int &k,const int &key)//插入结点
{
    if(!k)
    {
        k=++tot;
        v(k)=key;
        p(k)=rand();//p(k)为随机的优先级
        c(k)=s(k)=1;//新建结点
        lc(k)=rc(k)=0;
        return ;
    }
    else
    ++s(k);
    if(v(k)==key)//重复元素处理
    ++c(k);
    else
    if(key<v(k))
    {
        insert(lc(k),key);
        if(p(lc(k))<p(k))//维护Treap堆性质
        Zig(k);
    }
    else
    {
        insert(rc(k),key);
        if(p(rc(k))<p(k))
        Zag(k);//维护Treap堆性质
    }
    return ;
}
inline void Delete(int &k,const int &key)
{
    if(v(k)==key)
    {
        if(c(k)>1)//重复元素
        --c(k),--s(k);
        else
        if(!lc(k)||!rc(k))//链结点可以直接删
        k=lc(k)+rc(k);
        else
        if(p(lc(k))<p(rc(k)))
        Zig(k),Delete(k,key);
        else
        Zag(k),Delete(k,key);
        return ;
    }
    --s(k);//维护size
    if(key<v(k))
    Delete(lc(k),key);
    else
    Delete(rc(k),key);
    return ;
}
inline int QueryPre(const int &key)//求一个元素在Treap中的前驱
{
    int x=rt,res=-inf;
    while(x)
    {
        if(v(x)<key)//要求的前驱为结点x或在结点x的右子树内
        res=v(x),x=rc(x);
        else
        x=lc(x);
    }
    return res;
}
inline int QuerySuf(const int &key)//求一个元素在Treap中的后继
{
    int x=rt,res=inf;
    while(x)
    {
        if(v(x)>key)//要求的前驱为结点x或在结点x的左子树内
        res=v(x),x=lc(x);
        else
        x=rc(x);
    }
    return res;
}
inline int QueryKth(int k)//在Treap中查找排名第k的元素
{
    int x=rt;
    while(x)
    {
        if(s(lc(x))<k&&s(lc(x))+c(x)>=k)//排名为k的元素在结点x
        return v(x);
        if(s(lc(x))>=k)//排名为k的元素在左子树
        x=lc(x);
        else//排名为k的元素在右子树
        k-=s(lc(x))+c(x),x=rc(x);
    }
    return 0;
}
inline int QueryRank(const int &key)//在Treap中求元素的排名
{
    int x=rt,res=0;
    while(x)
    {
        if(key==v(x))
        return res+s(lc(x))+1;//找到该元素并求得排名
        if(key<v(x))
        x=lc(x);
        else
        res+=s(lc(x))+c(x),x=rc(x);
    }
}
inline int get()
{
    char c;
    int sign=1,res;
    while((c=getchar())<'0'||c>'9')
    if(c=='-')
    sign=-1;
    res=c-'0';
    while((c=getchar())>='0'&&c<='9')
    res=res*10+c-'0';
    return res*sign;
}
int n;
int main()
{
    n=get();
    //cout<<n<<endl;
    int opt,x;
    while(n--)
    {
        opt=get();
        x=get();
        //cout<<opt<<" "<<x<<endl;
        if(opt==1)
        insert(rt,x);
        else
        if(opt==2)
        Delete(rt,x);
        else
        if(opt==3)
        printf("%d\n",QueryRank(x));
        else
        if(opt==4)
        printf("%d\n",QueryKth(x));
        else
        if(opt==5)
        printf("%d\n",QueryPre(x));
        else
        printf("%d\n",QuerySuf(x));
    }
    return 0;
}