luogu3369 模板普通平衡树（Treap/SBT）

Posted 2020-10-16 poorpool

treap做法，参考hzwer的博客

#include <iostream>
#include <cstdlib>
#include <cstdio>
using namespace std;
struct Node{
    int val, l, r, sze, rnd, hav;
    //val代表当前节点的值，l、r是左右孩子编号，sze是以当前节点为根的子树的大小
    //rnd是一个随机出来的权值， hav是当前节点的值对应的数字有几个
}nd[100005];
int n, opt, x, rot, ans, num;
void update(int k){
    nd[k].sze = nd[nd[k].l].sze + nd[nd[k].r].sze + nd[k].hav;
}
void lRotate(int &k){
    int t=nd[k].r; nd[k].r = nd[t].l; nd[t].l = k;
    nd[t].sze = nd[k].sze; update(k); k = t;
}
void rRotate(int &k){
    int t=nd[k].l; nd[k].l = nd[t].r; nd[t].r = k;
    nd[t].sze = nd[k].sze; update(k); k = t;
}
void ins(int &k, int x){
    if(!k){
        k = ++num; nd[k].sze = nd[k].hav = 1;
        nd[k].val = x; nd[k].rnd = rand();
        //创建一个新节点
        return ;
    }
    nd[k].sze++;
    if(nd[k].val==x)    nd[k].hav++;
    else if(x>nd[k].val){
        ins(nd[k].r, x);
        if(nd[nd[k].r].rnd<nd[k].rnd)   lRotate(k);
        //不满足堆的性质则旋转
    }
    else{
        ins(nd[k].l, x);
        if(nd[nd[k].l].rnd<nd[k].rnd)   rRotate(k);
        //不满足堆的性质则旋转
    }
}
void del(int &k, int x){
    if(!k)  return ;
    if(nd[k].val==x){
        if(nd[k].hav>1){
            nd[k].hav--; nd[k].sze--;
            //要是有不止一个的话减一就好了
            return ;
        }
        if(nd[k].l*nd[k].r==0)  k = nd[k].l + nd[k].r;
        //如果左右子树有一个为空，则把k置为不为空的那颗子树
        //如果均为空，那么k也变成了空
        else if(nd[nd[k].l].rnd<nd[nd[k].r].rnd)
            rRotate(k), del(k, x);
        else    lRotate(k), del(k, x);
    }
    else if(x>nd[k].val)    nd[k].sze--, del(nd[k].r, x);
    else    nd[k].sze--, del(nd[k].l, x);
}
int queryRank(int k, int x){
    if(!k)  return 0;
    if(nd[k].val==x)    return nd[nd[k].l].sze+1;
    else if(nd[k].val<x)    return nd[nd[k].l].sze+nd[k].hav+queryRank(nd[k].r, x);
    else return queryRank(nd[k].l, x);
}
int queryNum(int k, int x){
    if(!k)  return 0;
    if(x<=nd[nd[k].l].sze)  return queryNum(nd[k].l, x);
    else if(x>nd[nd[k].l].sze+nd[k].hav)    return queryNum(nd[k].r, x-nd[nd[k].l].sze-nd[k].hav);
    else return nd[k].val;
}
void queryPre(int k, int x){
    if(!k)  return ;
    if(nd[k].val<x) ans = k, queryPre(nd[k].r, x);
    else queryPre(nd[k].l, x);
}
void queryNxt(int k, int x){
    if(!k)  return ;
    if(nd[k].val>x) ans = k, queryNxt(nd[k].l, x);
    else queryNxt(nd[k].r, x);
}
int main(){
    cin>>n;
    while(n--){
        scanf("%d %d", &opt, &x);
        if(opt==1)  ins(rot, x);
        if(opt==2)  del(rot, x);
        if(opt==3)  printf("%d\n", queryRank(rot, x));
        if(opt==4)  printf("%d\n", queryNum(rot, x));
        if(opt==5)  queryPre(rot, x), printf("%d\n", nd[ans].val);
        if(opt==6)  queryNxt(rot, x), printf("%d\n", nd[ans].val);
    }
    return 0;
}