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

BZOJ传送门
平衡树模板题，注意相同元素的处理。
Treap：

#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;    
}