|BZOJ 3224|平衡树|Tyvj 1728 普通平衡树

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<vector>
#define ms(i,j) memset(i,j, sizeof i);
using namespace std;
struct node
{
node *ch[2];//左右孩子
int v, r;//值，优先级
int s;//附加值：以当前节点为根的结点数量
int w;//附加值：和当前节点相同值的结点数
void mt()
{
s = w;
if (ch[0]!=NULL) s += ch[0]->s;
if (ch[1]!=NULL) s += ch[1]->s;
}
};
int ans;
struct treap
{
node *root;
void rotate(int d, node *&o)//d=0则左旋 d=1则右旋
{
node *k = o->ch[d^1]; o->ch[d^1] = k->ch[d];
k->ch[d] = o; o->mt(); k->mt(); o = k;
}
void insert(int x, node *&o)//插入一个数
{
if (o==NULL)
{
o = new node(); o->v = x; o->r = rand(); o->s = o->w = 1; o->ch[0] = o->ch[1] = NULL; //初值
} else  if (o->v==x) o->w++; //有相同直接w++
else
{
int d = (x < o->v ? 0 : 1);
insert(x, o->ch[d]);
if (o->ch[d]->r > o->r) rotate(d^1, o);
}
o->mt();
}
void del(int x, node *&o)//删除一个数
{
int d = (x < o->v ? 0 : 1);
if (o->v==x)//找到
{
if (o->w>1) {o->w--; o->s--;} else//不止一个数就直接w--，s--
if (o->ch[0]==NULL) o = o->ch[1];
else if (o->ch[1]==NULL) o = o->ch[0];
else {
int d2 = (o->ch[0]->r > o->ch[1]->r ? 1 : 0);
rotate(d2,o); del(x, o->ch[d2]);
}
} else
{
del(x, o->ch[d]);
}
if (o!=NULL) o->mt();
}
int rank(int x, node *o)//求x的排名
{
int tmp;
if (o->ch[0]==NULL) tmp = 0;
else  tmp=o->ch[0]->s;//求s

if (o->v==x) return tmp+1;//找到了
else if (o->v >x) return rank(x,o->ch[0]);
else return tmp+o->w+rank(x,o->ch[1]);
}
int kth(int k, node *o)//求第k小
{
if (o==NULL||o->s<k||k<=0) return -1;//不符合要求
int tmp;
if (o->ch[0]==NULL) tmp = 0;
else tmp = o->ch[0]->s;//求s

if (k<=tmp) return kth(k,o->ch[0]);
else if (k > tmp+o->w) return kth(k - tmp - o->w,o->ch[1]);
else return o->v;//找到了
}
void pred(int x, node *o)//求前驱
{
if(o==NULL)return;
if(o->v<x) {
ans = o->v;
pred(x, o->ch[1]);
} else pred(x, o->ch[0]);
}
void succ(int x, node *o)//求后继
{
if(o==NULL)return;
if(o->v>x) {
ans = o->v;
succ(x, o->ch[0]);
} else succ(x, o->ch[1]);
}
};
treap tree;
int n;
int main()
{
scanf("%d", &n);
int opt,x;
for (int i=1;i<=n;i++)
{
scanf("%d%d", &opt, &x);
switch(opt)
{
case 1: tree.insert(x, tree.root); break;
case 2: tree.del(x, tree.root); break;
case 3: printf("%d\n", tree.rank(x, tree.root)); break;
case 4: printf("%d\n", tree.kth(x, tree.root)); break;
case 5: tree.pred(x,tree.root); printf("%d\n", ans); break;
case 6: tree.succ(x,tree.root); printf("%d\n", ans); break;
}
}
return 0;
}