Treap的实现方法 [BZOJ 3224]

传说，有一种排序二叉树叫做Treap。
而 Treap = Tree + Heap
所以，Treap既具有树，也具有堆的性质。
它的基本操作和普通的树相近，但也有一些差异。

（以上全部为乱讲系列）

（如果要看详细介绍，这里给出lmy大神关于平衡树的研究讲解http://blog.csdn.net/lemonoil/article/details/54405613)

---

（详细介绍之后会补充的）

---

如何写出一个合格的Treap？

首先，它的存储方式和其他的二叉树类似，都有关于该点以及它的子树的存储信息。我们用 val 保存每个点的点权，rnd (rd) 保存每个点拥有的随机优先级，这样我们就可以在之后利用堆的性质对其加以排序。同时我们也要记录这个树所包含的所有子树以及它自己的大小，用 size 来表示。如果一个点可以被多次建立，那么我们还需要用 weight (wei) 来保存当前点有多少是“重复”的 （即：这个点被重复插入了多少次）。

struct Treap {
    int val,rd,size,wei;
    Treap *ch[2];
    Treap (int x){
        this->val = x;    // this 在这里可以省略
        rd = rand();      // rand() 可以优化，后面有介绍
        size = wei = 1;
        ch[0] = ch[1] = NULL;
    }
    int comp (int x){
        if (x == val) return -1;
        return x < val ? 0 : 1;
    }
    void maintain (){
        size = this->wei;   // this 同上
        if (s[0] != NULL) size += s[0]->size;
        if (s[1] != NULL) size += s[1]->size;
    }
};

Treap *root = NULL;

接下来是它的两个基本操作，插入和删除。

那么问题来了：如何对其进行旋转操作呢？这里我们给出了一种方法。至于实现图例，本人还没有添加，先看代码。

void rotate(Treap *&t, int d){    //旋转操作
    Treap *m = t->s[d ^ 1];
    t->s[d ^ 1] = m->s[d];
    m->s[d] = t;
                                 //这里的旋转操作很重要，需要仔细理解
    t->maintain();               //必须先维护t，t此时是m的子树
    m->maintain();
    t = m;
}
void ins(Treap *&t, int v){
    if (t == NULL){            //当前节点不存在
        t = new Treap(v);
    } else {
        int d = t->comp(v);    //比较插入的值与这个点的大小；
                               //如果一样大就增加当前点的 size 和 weight ，否则往下继续找
        if (d == -1){
            t->size++;
            t->wei++;
        } else {
            t->size++;
            ins(t->s[d], v);
            if (t->s[d]->rd > t->rd){
                rotate(t, d^1);
            }
        }
        t->maintain();
    }
}
void del(Treap *&t,int v){
    int d = t->comp(v);
    if (d == -1){
        if (t->wei > 1){
            t->wei--;
            t->size--;
        } else {
            Treap *m = t;
            if (t->s[0] == NULL){        //左儿子为空 
                t = t->s[1];
                delete m;
                m = NULL;
            } else if (t->s[1] == NULL){ //右儿子为空 
                t = t->s[0];
                delete m;
                m = NULL;
            } else {             //两个儿子都保存有数据 
                int kd;
                kd = t->s[0]->rd < t->s[1]->rd ? 0 : 1;
                rotate(t,kd);
                del(t->s[kd], v);
            }
        }
    } else {
        t->size--;
        del(t->s[d], v);
    }
    if (t != NULL)                //防止删空
        t->maintain();
}

那么，这里的基本修改操作已经介绍完了。

所以，如何查询Treap里面的某些特定元素（某一个数的排名，处于某一个排名的数，某个特定的数的前驱/后继）？

这里我们就要用到堆的性质。
我们以查询一个数 k 的排名为例子，当 k 的值大于当前节点时，即搜索它的右子节点，同时加上该节点和它的左子树的 size 值（表示它的名次）。当 k 的值小于当前节点时，即搜索它的左子树，此时不加任何 size 值（k 的排名会在前面）。当 k 的值与当前节点相等时，返回 rank+1（名次要包括当前节点）。
其他的查询方式与上述思想相似，这里不再赘述。看代码就能懂了。

int query_rank(Treap *t, int v){       //查询一个数的排名
    int d = t->comp(v),rank;
    if (t->s[0] == NULL){
        rank = 0;
    }

//如果这个位置没有左子树，那么，这个数的相对名次（相对于这棵子树而言）为0，返回的名次值就应为1

    else {
        rank = t->s[0]->size;
    }

//返回所有左子树的size值，名次也对应地增加这么多

    if (d == -1){
        return rank + 1;
    } else if (d == 1){
        return query_rank(t->s[1], v) + t->wei + rank;
    } else {
        return query_rank(t->s[0], v);
    }
}

int query_kth(Treap *t, int k){       //查询排名为k的数
    if (t->s[0] == NULL){
        if (k >= 1 && k <= t->wei)
            return t->val;
        return query_kth(t->s[1], k-(t->wei) );
    } else {
        int p = t->s[0]->size + t->wei;
        if (k > p){
            return query_kth(t->s[1], k - p);
        } else if (k > t->s[0]->size && k <= p){
            return t->val;
        } else {
            return query_kth(t->s[0], k);
        }
    }
}

int query_prev(Treap *t, int v, int m){       //查询前驱
    if (t == NULL) return m;
    if (v > t->val){
        m = t->val;
        return query_prev(t->s[1], v, m);
    } else {
        return query_prev(t->s[0], v, m);
    }
}

int query_next(Treap *t, int v, int m){       //查询后继
    if (t == NULL) return m;
    if (v < t->val){
        maxt = t->val;
        return query_next(t->s[0], v, m);
    } else {
        return query_next(t->s[1], v, m);
    }
}

至此，我们已经明白了所有的操作方式。
下面贴上完整的代码：

#include<cstdio>
#include<iostream>
using namespace std;
inline int random();

struct Treap{
    int val,rd,size,wei;
    Treap *s[2];
    Treap(int a){
        this->val=a;
        size=1;wei=1;
        rd=random();
        s[0]=s[1]=NULL;
    }
    int comp(int a){
        if (a==val) return -1;
        return a<val ? 0:1;
    }
    void maintain(){
        size=this->wei;
        if (s[0]!=NULL) size+=s[0]->size;
        if (s[1]!=NULL) size+=s[1]->size;
    }
};
int n,opt,ap;

inline int random(){
    static int seed=703;
    return seed=(seed*48271LL%2147483647);
}
inline void rotate(Treap *&t,int d){
    Treap *m=t->s[d^1];
    t->s[d^1]=m->s[d];
    m->s[d]=t;
    t->maintain();
    m->maintain();
    t=m;
}
inline void ins(Treap *&t,int v){
    if (t==NULL){
        t=new Treap(v);
    } else {
        int d=t->comp(v);
        if (d==-1){
            t->size++;
            t->wei++;
        } else {
            t->size++;
            ins(t->s[d],v);
            if (t->s[d]->rd > t->rd){
                rotate(t,d^1);
            }
        }
        t->maintain();
    }
}
void del(Treap *&t,int v){
    int d=t->comp(v);
    if (d==-1){
        if (t->wei > 1){
            t->wei--;
            t->size--;
        } else {
            Treap *m=t;
            if (t->s[0]==NULL){//左儿子为空 
                t=t->s[1];
                delete m;
                m=NULL;
            } else if (t->s[1]==NULL){//右儿子为空 
                t=t->s[0];
                delete m;
                m=NULL;
            } else {//两个儿子都保存有数据 
                int kd;
                kd=t->s[0]->rd < t->s[1]->rd ? 0:1;
                rotate(t,kd);
                del(t->s[kd],v);
            }
        }
    } else {
        t->size--;
        del(t->s[d],v);
    }
    if (t!=NULL)
        t->maintain();
}
inline int query(Treap *&t,int v){
    int d=t->comp(v),rank;
    if (t->s[0]==NULL){
        rank=0;
    } else {
        rank=t->s[0]->size;
    }

    if (d==-1){
        return rank+1;
    } else if (d==1){
        return query(t->s[1],v) + t->wei + rank;
    } else {
        return query(t->s[0],v);
    }
}
inline int query_rank(Treap *&t,int k){
    if (t->s[0]==NULL){
        if (k>=1 && k <= t->wei)
            return t->val;
        return query_rank(t->s[1], k-(t->wei) );
    } else {
        if (k > t->s[0]->size + t->wei){
            return query_rank(t->s[1], k-(t->wei)-(t->s[0]->size) );
        } else if (k > t->s[0]->size && k <= t->s[0]->size + t->wei){
            return t->val;
        } else {
            return query_rank(t->s[0],k);
        }
    }
}
inline int query_prev(Treap *t,int v,int mint){
    if (t==NULL) return mint;
    if (v > t->val){
        mint=t->val;
        return query_prev(t->s[1],v,mint);
    } else {
        return query_prev(t->s[0],v,mint);
    }
}
inline int query_next(Treap *t,int v,int maxt){
    if (t==NULL) return maxt;
    if (v < t->val){
        maxt=t->val;
        return query_next(t->s[0],v,maxt);
    } else {
        return query_next(t->s[1],v,maxt);
    }
}
inline int read()
{
    int data=0,w=1; char ch=0;
    while(ch!='-' && (ch<'0' || ch>'9')) ch=getchar();
    if(ch=='-') w=-1,ch=getchar();
    while(ch>='0' && ch<='9') data=data*10+ch-'0',ch=getchar();
    return data*w;
}
int main(){
    Treap *root=NULL;
    n=read();
    for (register int i=1;i<=n;i++){
        opt=read();
        ap=read();
        switch (opt){
            case 1:ins(root,ap);break;
            case 2:del(root,ap);break;
            case 3:printf("%d\n",query(root,ap));break;
            case 4:printf("%d\n",query_rank(root,ap));break;
            case 5:printf("%d\n",query_prev(root,ap,0));break;
            case 6:printf("%d\n",query_next(root,ap,0));break;
        }
    }
    return 0;
}