在 P3369 【模板】普通平衡树 中的测试样例中,最后一个输出是正确的,除此之外的输出都是 inf。
代码有详细注释。
Right Output
106465
84185
492737
------------
My Output
-10000002
10000002
192737
//Treap 小根堆
#include<bits/stdc++.h>
using namespace std;
const int inf=1e7+2,N=1e5+5;
int n,idx=0,x,opt,root=0;
struct node{
int num; //数字
int val; //权值 (随机分配)
int l,r; //左右子树
int cnt; //这个数字的个数
int size; //自己加上子树的节点个数
}tr[N];
void pushup(int p){ //和线段树的操作一样,更新该节点的 size
tr[p].size=tr[tr[p].l].size + tr[tr[p].r].size + tr[p].cnt;
}
int new_node(int x){ //创建新节点
tr[++idx].num=x;
tr[idx].val=rand(); //权值随机
tr[idx].cnt= tr[idx].size=1; //叶子节点,大小为 1
return idx; //返回序号
}
void left(int &p){ //左旋,提右孩子
int y=tr[p].r;
tr[p].r=tr[y].l;
tr[y].l=p;
p=y; //实时更新根节点
pushup(tr[p].l);pushup(p); //重新统计
}
void right(int &p){ //与左旋相反 提升左孩子
int y=tr[p].l;
tr[p].l=tr[y].r;
tr[y].r=p;
p=y;
pushup(tr[p].r);pushup(p);
}
void build(){ //初始化
new_node(-inf);new_node(inf); //加入极值,防止越界
root=1; //根节点为 1
tr[1].r=2;
pushup(root);
if(tr[1].val>tr[2].val){ //如果权值不满足
left(root);
}
}
void insert(int &p,int x){ //插入节点
if(!p) p=new_node(x); //如果为空,新建一个节点
else if(tr[p].num==x) tr[p].cnt++; //如果值相同,有重复,数量加 1
else if(x>tr[p].num){ //插在右子树
insert(tr[p].r,x);
if(tr[tr[p].r].val<tr[p].val){ //维护小根堆性质
left(p); //转起来!
}
}
else{ //插在左子树
insert(tr[p].r,x);
if(tr[tr[p].r].val>tr[p].val){
right(p); //转起来!
}
}
pushup(p); //记得更新
}
void del(int &p,int x){ //删除
if(!p) return; //不存在,不处理
if(tr[p].num==x){ //若数字相同
if(tr[p].num>1) tr[p].cnt--; //如果数量大于 1,直接减数量
else if(tr[p].l||tr[p].r){ //若不是叶子节点
if(tr[p].r==0 || tr[tr[p].l].val < tr[tr[p].r].val){
//如果右节点为空 或 左节点小于右节点
right(p); //右 转起来!
del(tr[p].r,x); //继续递归
}
else{
left(p); //左 转起来!
del(tr[p].l,x); //继续递归
}
}
else{
p=0;
}
}
else if(x>tr[p].num){ //查找右子树
del(tr[p].r,x);
}
else{ //查找左子树
del(tr[p].l,x);
}
pushup(p); //记得更新
}
int num_to_rank(int p,int x){ //查询 x 的排名
if(!p) return 0; //如果不存在
if(tr[p].num==x){ //如果数值相等
return tr[tr[p].l].size+1; //比 x 小的就有 左子树大小加上 1
}
if(x<tr[p].num){ //否则 若在左子树
return num_to_rank(tr[p].l,x); //继续递归
}
return tr[tr[p].l].size + tr[p].cnt + num_to_rank(tr[p].r,x);
//否则在右子树 比 x 小的 左子树大小 + 根节点 + 递归右子树
}
int rank_to_num(int p,int rank){ //查询排名为 x 的数
if(!p) return inf; //如果不存在
if(tr[tr[p].l].size>=rank){ //如果数值相等
return num_to_rank(tr[p].l,rank); //rank 在左子树范围
}
if(x<=tr[tr[p].l].size + tr[p].cnt){ // x 在左子树 到 根节点的范围
return tr[p].num; //找到答案
}
return rank_to_num(tr[p].r,rank - tr[tr[p].l].size - tr[p].cnt);
//否则在右子树 因为rank是在整棵树中的,所以要rank需更新为 rank-左子树大小 + 根节点
}
int pre(int p,int x){ //查 小于 x 的最大数
if(!p) return -inf;
if(x <= tr[p].num){ //如果 在左子树
return pre(tr[p].l,x);
}
return max(tr[p].num , pre(tr[p].r,x));
//答案一定在 根节点 或右子树 返回二者的最大值
}
int after(int p,int x){ //查 大于 x 的最小数
if(!p) return inf;
if(x >= tr[p].num){ //如果 在右子树
return pre(tr[p].r,x);
}
return min(tr[p].num , pre(tr[p].l,x));
//答案一定在 根节点 或左子树 返回二者的最小值
}
int main(){
//ios::sync_with_stdio(0);
//cin.tie(0);cout.tie(0);
build();
cin>>n;
while(n--){
cin>>opt>>x;
if(opt==1) insert(root,x); //插入
if(opt==2) del(root,x); //删除
if(opt==3) cout<<num_to_rank(root,x)-1<<endl; //查询 x 的排名
if(opt==4) cout<<rank_to_num(root,x+1)<<endl; //查询第 x 名的数
if(opt==5) cout<<pre(root,x)<<endl; //查 小于 x 的最大数
if(opt==6) cout<<after(root,x)<<endl; //查 大于 x 的最小数
}
return 0;
}