洛谷 P3369 【模板】普通平衡树 仿刘汝佳风格代码的treap

P3369

Treap基本操作，当做模板题练练手了，看了好多题解代码，发现没有刘汝佳风格的代码，就决定写一份仿刘汝佳代码。

#include<cstdio>
#include<cstdlib> 
#include<algorithm>
using namespace std;
const int inf=1e9;
struct Node *null;
struct Node
{
	Node *ch[2];
	int r;
	int v;
	int w;//多出来的w是记录相同的数的个数 
	int s;
	Node(int v):v(v){ch[0]=ch[1]=null;r=rand();s=w=1;}
	bool operator<(const Node& rhs)const{
		return r<rhs.r;
	}
	int cmp(int x)const{
		if(x==v)return -1;
		return x<v?0:1;
	}
	void maintain()
	{
		s=w;
		s+=ch[0]->s;
		s+=ch[1]->s;
	}
};
void rotate(Node* &o,int d)
{
	Node *k=o->ch[d^1];o->ch[d^1]=k->ch[d];
	k->ch[d]=o;o->maintain();k->maintain();
	o=k;
}
void insert(Node* &o,int x)
{
	if(o==null) o=new Node(x);
	else if(o->v==x)
	o->w++,o->maintain();
	else
	{
		int d=(x<o->v?0:1);
		insert(o->ch[d],x);
		if(o->ch[d]->r>o->r)rotate(o,d^1);
	}
	o->maintain();
}
void remove(Node* &o,int x){
	int d=o->cmp(x);
	if(d==-1){
		if(o->w>1)
		{
			o->w--;
			o->maintain();
			return;
		}
		Node* u=o;
		if(o->ch[0]!=null&&o->ch[1]!=null)
		{
			int d2=(o->ch[0]->r>o->ch[1]->r?1:0);
			rotate(o,d2);remove(o->ch[d2],x);
		}else{
			if(o->ch[0]==null)
			o=o->ch[1];
			else o=o->ch[0];
			delete u;
		}
	}
	else
	remove(o->ch[d],x);
	if(o!=null)o->maintain();
}
int kth(Node* o,int k)
{
	if(o==null||k<=0||k>o->s)return 0;
	int s=(o->ch[0]==null?0:o->ch[0]->s);
	if(k>s&&k<=s+o->w)return o->v;
	else if(k<=s)return kth(o->ch[0],k);
	else return kth(o->ch[1],k-s-o->w);
}
int rank(Node* o,int k,int sum)
{
	if(o==null)return sum+1;
	if(o->v==k)return o->ch[0]->s+1+sum;
	else if(k<o->v)return rank(o->ch[0],k,sum);
	else return rank(o->ch[1],k,sum+o->ch[0]->s+o->w);
}
int ans;
void pre(Node *o,int k)
{
	if(o==null)return;
	if(o->v<k)
	{
		ans=max(ans,o->v);
		pre(o->ch[1],k);
	}
	else
	pre(o->ch[0],k);	
}
void next(Node *o,int k)
{
	if(o==null)return;
	if(o->v>k)
	{
		ans=min(ans,o->v);
		next(o->ch[0],k); 
	}
	else
	next(o->ch[1],k);
}
int main()
{
	int q,op,x;
	scanf("%d",&q);
	null=new Node(0);
	null->s=null->w=0;
	Node *root=null;
	while(q--) 
	{
		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",rank(root,x,0));
		else if(op==4) printf("%d\n",kth(root,x));
		else if(op==5) ans=-inf,pre(root,x),printf("%d\n",ans);
		else ans=inf,next(root,x),printf("%d\n",ans);
	}
	
}