BZOJ 3224 普通平衡树

BZOJ 3224

懒羊羊看了下BST treap 

然后对着敲一下BZOJ 3224

代码有注释 额 感觉treap好玄学啊 rand是什么鬼

#include <cstdio>
#include <iostream>
#include <queue>
#include <cmath>
#include <cstring>
#include <algorithm>
#include <stack>
#include <map>
#include <set>
#include <sstream>
using namespace std;
const int INF = 0x7fffffff;
int tot,root,n;
const int SIZE = 100010;
struct Treap{
    int l,r;
    int val,dat;
    int cnt,size;
}a[SIZE];
 
int New(int val){
    a[++tot].val = val;
    a[tot].dat = rand();
    a[tot].cnt = a[tot].size = 1;
    return tot;
}
 
void Update(int p){
    a[p].size = a[a[p].l].size + a[a[p].r].size + a[p].cnt;
}
 
void Build(){
    New(-INF),New(INF);
    root = 1,a[1].r = 2;
    Update(root);
}
 
int Getrankbyval(int p,int val){
    if(p==0) return 0;
    if(val==a[p].val) return a[a[p].l].size + 1;
    if(val<a[p].val) return Getrankbyval(a[p].l,val);
    return Getrankbyval(a[p].r,val) + a[a[p].l].size + a[p].cnt;
}
 
int Getvalbyrank(int p,int Rank){
    if(p==0) return INF;
    if(a[a[p].l].size>=Rank) return Getvalbyrank(a[p].l,Rank);
    if(a[a[p].l].size + a[p].cnt >= Rank) return a[p].val;
    return Getvalbyrank(a[p].r,Rank-a[a[p].l].size - a[p].cnt);
}
 
void zig(int &p){//右旋
    int q = a[p].l;
    a[p].l = a[q].r, a[q].r = p,p = q;
    Update(a[p].r),Update(p);
}
 
void zag(int &p){//左旋
    int q = a[p].r;
    a[p].r = a[q].l,a[q].l = p,p = q;
    Update(a[p].l),Update(p);
}
 
void Insert(int &p,int val){
    if(p==0){
        p = New(val);
        return ;
    }
    if(val == a[p].val){
        a[p].cnt++,Update(p);
        return ;
    }
    if(val<a[p].val){
        Insert(a[p].l,val);
        if(a[p].dat < a[a[p].l].dat) zig(p);
    }
    else {
        Insert(a[p].r,val);
        if(a[p].dat < a[a[p].r].dat) zag(p);
    }
    Update(p);
}
 
int Getpre(int val){
    int ans = 1;//a[1].val = -INF
    int p = root;
    while(p){
        if(val == a[p].val){
        if(a[p].l>0){
            p = a[p].l;
            while(a[p].r>0) p = a[p].r;
            ans = p;
        }
        break;
    }
    if(a[p].val < val && a[p].val > a[ans].val) ans = p;
    p = val < a[p].val ? a[p].l : a[p].r;
    }
    return a[ans].val;
}
 
int Getnext(int val){
    int ans = 2;//a[2].val = INF;
    int p = root;
    while(p){
        if(val==a[p].val){
            if(a[p].r>0){
                p = a[p].r;
                while(a[p].l>0) p = a[p].l;
                ans = p;
            }
            break;
        }
        if(a[p].val > val && a[p].val < a[ans].val) ans=  p;
        p = val < a[p].val ? a[p].l : a[p].r;
    }
    return a[ans].val;
}
 
void Remove(int &p,int val){
    if(p==0) return;
    if(val==a[p].val){
        if(a[p].cnt>1){
            a[p].cnt--,Update(p);
            return;
        }
        if(a[p].l||a[p].r){
            if(a[p].r==0||a[a[p].l].dat > a[a[p].r].dat)
                zig(p),Remove(a[p].r,val);
            else
                zag(p),Remove(a[p].l,val);
            Update(p);
        }
        else p = 0;
        return;
    }
    val < a[p].val ? Remove(a[p].l,val) : Remove(a[p].r,val);
    Update(p);
}
 
int main(){
    Build();
    scanf("%d",&n);
    while(n--){
        int opt,x;
        scanf("%d%d",&opt,&x);
        switch(opt){
        case 1:
            Insert(root,x);
            break;
        case 2:
            Remove(root,x);
            break;
        case 3:
            printf("%d\n",Getrankbyval(root,x) - 1);
            break;
        case 4:
            printf("%d\n",Getvalbyrank(root,x+1));
            break;
        case 5:
            printf("%d\n",Getpre(x));
            break;
        case 6:
            printf("%d\n",Getnext(x));
            break;
        }
    }
    return 0;
}