普通平衡树（Treap模板题）

题目描述

您需要写一种数据结构（可参考题目标题），来维护一些数，其中需要提供以下操作：

1. 插入数值x。
2. 删除数值x(若有多个相同的数，应只删除一个)。
3. 查询数值x的排名(若有多个相同的数，应输出最小的排名)。
4. 查询排名为x的数值。
5. 求数值x的前驱(前驱定义为小于x的最大的数)。
6. 求数值x的后继(后继定义为大于x的最小的数)。

注意： 数据保证查询的结果一定存在。

输入格式

第一行为n，表示操作的个数。
接下来n行每行有两个数opt和x，opt表示操作的序号(1<=opt<=6)。

输出格式

对于操作3,4,5,6每行输出一个数，表示对应答案。

数据范围

n≤100000,所有数均在−107到107内。

输入样例
8
1 10
1 20
1 30
3 20
4 2
2 10
5 25
6 -1
输出样例
2
20
20
20

题目分析

这就是一道平衡树（Treap）的模板题，题中要求的操作都是Treap的基础操作。

代码如下

#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#define LL long long
#define PII pair
#define L tr[u].l
#define R tr[u].r
using namespace std;
const int N=1e5+5,INF=1e8;
struct Node{
    int l,r;
    int key,val;		//key:Treap的关键字，Treap节点的位置是由key确定的
    int cnt,size;		//cnt记录这个节点中所含有的数；size以此节点为根的子树中包含的节点个数
}tr[N];
int root,idx;		//root记录根节点，idx记录当前创建的节点数量
void pushup(int u)	//由左右儿子的节点信息推出父节点的信息
{
    tr[u].size=tr[L].size+tr[R].size+tr[u].cnt;
}
int getNode(int key)		//创建新节点
{
    tr[++idx].key=key;
    tr[idx].val=rand();		//val要赋一个随机值，以保证树的平衡性
    tr[idx].cnt=tr[idx].size=1;	//初始化两个辅助变量（这两个一开始都是1）
    return idx;				//返回该节点的下标
}
void zig(int &u)    //右旋（模板）
{
    int q=L;
    L=tr[q].r,tr[q].r=u;
    u=q;
    pushup(R);
    pushup(u);
}
void zag(int &u)    //左旋（模板）
{
    int q=R;
    R=tr[q].l,tr[q].l=u;
    u=q;
    pushup(L);
    pushup(u);
}
void insert(int &u,int key)		//插入一个数
{
    if(!u) u=getNode(key);		//找到一个空位，放入该节点
    else if(tr[u].key==key) tr[u].cnt++;	//如果当前key的位置已经有数了，让cnt++即可
    else if(tr[u].key>key)		//如果当前节点的key大于所给的key，那么往左放
    {
        insert(L,key);
        if(tr[L].val>tr[u].val) zig(u);		//调整Treap的结构，使其趋于平衡
    }
    else		//如果当前节点的key小于所给的key，那么往右放
    {
        insert(R,key);
        if(tr[R].val>tr[u].val) zag(u);
    }
    pushup(u);		//更新该节点信息
}
void remove(int &u,int key)		//从树中删除一个节点
{
    if(!u) return;				//如果该节点不存在，则直接结束
    if(tr[u].key==key)
    {
        if(tr[u].cnt>1) tr[u].cnt--;	//如果该节点上的树大于1个，那么直接cnt--
        else if(L||R)			//查看该节点是否为叶节点
        {
            if(!R||tr[L].val>tr[R].val)	//不是则将该节点变为叶结点，再删除
            {
                zig(u);
                remove(R,key);
            }
            else
            {
                zag(u);
                remove(L,key);
            }
        }
        else u=0;		//是叶结点则直接删除
    }
    else if(tr[u].key>key) remove(L,key);	//当前节点的key大于所给key，往左找
    else remove(R,key);		//否则往右找
    if(u) pushup(u);		//如果该节点未被删除，那么更新该点
}
int kth(int &u,int key)		//返回这个key在这棵树的排名（第几小的数）
{
    if(!u) return 0;		//未找到该节点返回0
    if(tr[u].key==key) return tr[L].size+1;		//找到了该节点，该节点的排名即为左子树的大小+1
    if(tr[u].key>key) return kth(L,key);		//如果当前节点的key大于所给出的key，往左找
    return tr[L].size+tr[u].cnt+kth(R,key);		//否则往右找
}
int find(int &u,int rank)		//找出树中第k小的数
{
    if(!u) return INF;			//树中不足k个数，返回INF
    if(tr[L].size>=rank) return find(L,rank);	//此时说明要找的点在左子树中，往左找
    if(tr[L].size+tr[u].cnt>=rank) return tr[u].key;	//说明要找的点就是该节点，返回该节点的key
    return find(R,rank-tr[L].size-tr[u].cnt);	//说明要找的点在右子树中，往右找
}
int getPrev(int u,int key)     //找到严格小于key的最大值（前驱）
{
    if(!u) return -INF;			//没有返回-INF
    if(tr[u].key>=key) return getPrev(L,key);	//当前节点大于key，不满足条件，往左找
    return max(tr[u].key,getPrev(R,key));	//当前节点小于key，满足条件，但再看看有没有更符合条件的
}
int getNext(int u,int key)     //找到严格大于key的最小值（后继）
{
    if(!u) return INF;			//原理与找前驱类似
    if(tr[u].key<=key) return getNext(R,key);
    return min(tr[u].key,getNext(L,key));
}
int main()
{
    int n;
    scanf("%d",&n);
    while(n--)
    {
        int op,x;
        scanf("%d %d",&op,&x);
        if(op==1) insert(root,x);			//插入一个数
        else if(op==2) remove(root,x);		//删除一个数
        else if(op==3) printf("%d\n",kth(root,x));		//输出x在树中的排名
        else if(op==4) printf("%d\n",find(root,x));		//输出树中第x小的数
        else if(op==5) printf("%d\n",getPrev(root,x));	//找到x的前驱
        else printf("%d\n",getNext(root,x));			//找到x的后继
    }
    return 0;
}