平衡树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;
}
更多精彩
