平衡树Treap

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<algorithm>
using namespace std;
const int INF=0x3fffffff;
int node_cnt;																//结点计数器
struct 	Node;
typedef Node * pNode; 														//结点指针
struct Node																	//结点定义
{
	int	  key, val, size, same;
	pNode ls, rs;
	Node(int val):val(val), size(1), same(1){ ls=rs=NULL; key=rand(); }
	int  lsize(){ return ls?ls->size:0; }									//避免空指针查询
	int  rsize(){ return rs?rs->size:0; }									//避免空指针查询
	void update(){ size=lsize()+rsize()+same; }
};
void zig(pNode &p)															//左旋
{
	pNode tmp=p->rs; p->rs=tmp->ls; tmp->ls=p; tmp->size=p->size; p->update(); p=tmp;
}
void zag(pNode &p)															//右旋
{
	pNode tmp=p->ls; p->ls=tmp->rs; tmp->rs=p; tmp->size=p->size; p->update(); p=tmp;
}
void treap_insert(int x, pNode &p)											//插入结点
{
	if(!p){ p=new Node(x); node_cnt++; return; };
	if(p->val==x)	{ p->same++; p->update(); node_cnt++; return; }
	if(x< p->val)	{ treap_insert(x, p->ls); if(p->key>p->ls->key) zag(p); }
	if(x> p->val)	{ treap_insert(x, p->rs); if(p->key>p->rs->key) zig(p); }
	p->update();
}
void treap_delete(int x, pNode &p)											//删除结点
{
	if(!p)	return;														 			//没找到
	if(x<p->val)	treap_delete(x, p->ls);
	if(x>p->val)	treap_delete(x, p->rs);
	if(x==p->val)
	{
		if(p->same>1)	{ p->same--;  p->update(); node_cnt--;	return; }			//多个相同的，删除一个
		if(!(p->ls) && !(p->rs)){ delete p; p=NULL; node_cnt--;	return; }			//没有子树，直接删除
		else if(!p->ls && p->rs){ delete p; p=p->rs; node_cnt--; 		}			//只有右子树
		else if(p->ls && !p->rs){ delete p; p=p->ls; node_cnt--;		}			//只有左子树
		else															 			//有左、右子树
		{
			if(p->ls->key<p->rs->key)	zag(p), treap_delete(x, p->rs);
			else						zig(p), treap_delete(x, p->ls);
		}
	}
	p->update();
}
int  treap_rank(int x, const pNode &p)										//查询元素排名
{
	if(!p)	return 0;																//没找到
	if(p->val==x)		return p->lsize()+1;
	else if(p->val <x)	return p->lsize()+p->same+treap_rank(x, p->rs);
	else				return treap_rank(x, p->ls); 								//这个函数是存在问题的，如果要查找的数字不存在，排名就是N或0，该如何解决？
}
int  treap_find(int k, const pNode &p)										//查询排名为k的元素
{
	if(k>node_cnt)	return INF;														//如果k大于实际节点数，返回无穷大
	int tmp=p->lsize()+p->same;
	if(p->lsize()<k && k<=tmp)	return p->val;
	else if(k>tmp)				return treap_find(k-tmp, p->rs);
	else						return treap_find(k, p->ls);
}
int  treap_pre(int x, const pNode &p)										//查找x的前驱
{
	if(!p)	return -INF;															//找不到返回无穷小
	if(x<=p->val)	return treap_pre(x, p->ls);
	else			return max(p->val, treap_pre(x, p->rs));
}
int  treap_next(int x, const pNode &p)										//查找x的后继
{
	if(!p)	return INF;																//找不到返回无穷大
	if(x>=p->val)	return treap_next(x, p->rs);
	else			return min(p->val, treap_next(x, p->ls));
}
int main()
{
	int   n;
	pNode root=NULL;
	scanf("%d", &n);
	for(int i=1, opt, x; i<=n; i++)
	{
		scanf("%d%d", &opt, &x);
		if(opt==1)	treap_insert(x, root);
		if(opt==2)	treap_delete(x, root);
		if(opt==3)	printf("%d\n", treap_rank(x, root));
		if(opt==4)	printf("%d\n", treap_find(x, root));
		if(opt==5)	printf("%d\n", treap_pre(x, root) );
		if(opt==6)	printf("%d\n", treap_next(x, root));
	}
	return 0;
}