题目链接: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;
}