学习笔记-平衡树

treap

#include<bits/stdc++.h>
#define int long long
using namespace std;
#define ls t[x].ch[0]
#define rs t[x].ch[1]
const int N=114514;
const int inf=2147483647;
int cnt=0,root;
mt19937 rnd(0x7f);
struct treap{
	int ch[2], cnt, size, val, rd;
	//treap不需要记录父指针，rd表示节点的随机值
}t[N];

void up(int x){
	t[x].size = t[t[x].ch[0]].size+t[t[x].ch[1]].size+t[x].cnt;
}

void rotate(int &x, int d){//x代表的是旋转时作为父节点的节点,d代表的是旋转的方向
//d==0时是左儿子旋上来, d==1是右儿子旋上来.
    int son = t[x].ch[d];
    t[x].ch[d] = t[son].ch[d^1];
    t[son].ch[d^1] = x; up(x), up(x=son);//相当于up(son)
}

void insert(int &x, int val){
    if(!x){//找到对应位置就新建节点
        x = ++cnt;
        t[x].cnt = t[x].size = 1;
        t[x].val = val, t[x].rd = rnd();
        return;
    }
    t[x].size++;//因为插入了数,所以在路径上每个节点的size都会加1
    if(t[x].val == val){t[x].cnt++; return;}//找到了直接返回
    int d = t[x].val < val; insert(t[x].ch[d], val);//否则递归查找插入位置
    if(t[x].rd > t[t[x].ch[d]].rd) rotate(x, d);
}

void delet(int &x, int val){
    if(!x) return;//防止越界
    if(t[x].val == val){
        if(t[x].cnt > 1){t[x].cnt--, t[x].size--;return;}//有相同的就直接cnt--
        bool d = t[ls].rd > t[rs].rd;
        if(ls == 0 || rs == 0) x = ls+rs;//只有一个儿子就直接把那个儿子放到这个位置
        else rotate(x, d), delet(x, val);//否则将x旋下去,找一个随机值小的替代,直到回到1,2种情况
    }
    else t[x].size--, delet(t[x].ch[t[x].val<val], val);//递归找到要删除的节点.
}

int rk(int x, int val){
    if(!x) return 0;
    if(t[x].val == val) return t[ls].size+1;//找到了就返回最小的那个
    if(t[x].val > val) return rk(ls, val);//如果查找的数在x的左边,则直接往左边查
    return rk(rs, val)+t[ls].size+t[x].cnt;//否则往右边查,左边的所有数累加进答案
}

int kth(int root, int k){
    int x = root;
    while(1){
        if(k <= t[ls].size) x = ls;
        else if(k > t[ls].size+t[x].cnt)
            k -= t[ls].size+t[x].cnt, x = rs;
		else return t[x].val;
    }
}

int pre(int x, int val){
    if(!x) return -inf;//防止越界,同时-inf无法更新答案,
    if(t[x].val >= val) return pre(ls, val);//如果该节点的权值大于等于要找的权值
    //则不能成为前驱,递归查找左子树(有可能找到前驱)
    return max(pre(rs, val), t[x].val);//找右子树中是否存在前驱
}

int nex(int x, int val){//同上
    if(!x) return inf;
    if(t[x].val <= val) return nex(rs, val);
    return min(nex(ls, val), t[x].val);
}

signed main(){
	#ifndef ONLINE_JUDGE
        freopen("1.in","r",stdin);
        freopen("1.out","w",stdout);
    #endif
	int m,op,x;
    scanf("%lld",&m);
    while(m--)
    {
        scanf("%lld%lld",&op,&x);
        switch(op)
        {
            case 1:{insert(root,x);break;}
            case 2:{delet(root,x);break;}
            case 3:{printf("%lld\n",rk(root,x));break;}
            case 4:{printf("%lld\n",kth(root,x));break;}
            case 5:{printf("%lld\n",pre(root,x));break;}
            case 6:{printf("%lld\n",nex(root,x));break;}
        }
    }
}

FHQ-treap

#include<bits/stdc++.h>
#define N 210000
#define llt long long
using namespace std;
llt root,m,op,x,tot; 
struct TREAP
{
    #define ls(now) tree[now].son[0]
    #define rs(now) tree[now].son[1]
    #define v(now)  tree[now].node
    #define upd(now)    tree[now].siz=tree[now].cnt+tree[ls(now)].siz+tree[rs(now)].siz
    struct NODE{llt node,son[2],rd,cnt,siz;}tree[N];llt sz;
    llt create(llt value){v(++sz)=value;tree[sz].cnt=tree[sz].siz=1;tree[sz].rd=rand();return sz; }
    void split(llt now,llt num,llt &x,llt &y)
    {
        if(!now) {x=y=0;return;}
        if(v(now)>num)  y=now,split(ls(now),num,x,ls(now));
        if(v(now)<=num) x=now,split(rs(now),num,rs(now),y);
        upd(now);
    }
    void merge(llt &now,llt x,llt y)
    {
        if(x==0)    {now=y;return;}
        if(y==0)    {now=x;return;}
        if(tree[x].rd>tree[y].rd)   merge(rs(x),rs(x),y),upd(now=x);
        else                        merge(ls(y),x,ls(y)),upd(now=y);
    }
    void insert(llt &now,llt nnd)
    {
        llt ltree,rtree,is;
        split(now,nnd-1,ltree,rtree);  
        merge(now,ltree,create(nnd)),merge(now,now,rtree);
    }
    void remove(llt &now,llt num)
    {
        llt ltree,rtree,is;
        split(now,num-1,ltree,rtree);split(rtree,num,is,rtree);
        merge(is,ls(is),rs(is));merge(now,ltree,is);merge(now,now,rtree);
    }
    llt check_rk(llt now,llt num)
    {
        llt ltree,rtree,ans;
        split(now,num-1,ltree,rtree);
        ans=tree[ltree].siz+1;
        merge(now,ltree,rtree);
        return ans;
    }
    llt check_num(llt now,llt rk)
    {
        while(1)
        {
            if(tree[ls(now)].siz>=rk)  now=ls(now);
            else if(tree[ls(now)].siz+tree[now].cnt>=rk)    break;
            else rk-=tree[ls(now)].siz+tree[now].cnt,now=rs(now);
        }
        return v(now);
    }
    llt Pre(llt &now,llt num)
    {
        llt ltree,rtree,is;
        split(now,num-1,ltree,rtree),is=ltree;
        while(rs(is))   is=rs(is);
        merge(now,ltree,rtree);
        return v(is);
    }
    llt Next(llt &now,llt num)
    {
        llt ltree,rtree,is;
        split(now,num,ltree,rtree),is=rtree;
        while(ls(is))   is=ls(is);
        merge(now,ltree,rtree);
        return v(is);
    }
}treap;
int main()
{
    #ifndef ONLINE_JUDGE
        freopen("1.in","r",stdin);
        freopen("1.out","w",stdout);
    #endif
    srand(time(0));
    scanf("%lld",&m);
    while(m--)
    {
        scanf("%lld%lld",&op,&x);
        switch(op)
        {
            case 1:{treap.insert(root,x);break;}
            case 2:{treap.remove(root,x);break;}
            case 3:{printf("%lld\n",treap.check_rk(root,x));break;}
            case 4:{printf("%lld\n",treap.check_num(root,x));break;}
            case 5:{printf("%lld\n",treap.Pre(root,x));break;}
            case 6:{printf("%lld\n",treap.Next(root,x));break;}
        }
    }
    return 0;
}