Treap模板

Posted 2020-10-27 xuanyi

\(treap\) （树堆）= \(tree\) （树） + \(heap\) （堆）

显然,这个名字一定是二叉排序树和堆的结合

事实就是这样

怎么定义?

• 是一棵二叉树,但并不一定是完全二叉树
• 按权值形成二叉排序树,即左边比"我"小,右边比"我"大
• 按随机给的值形成一棵堆,以小根堆为例,即"我"的孩子的随机值小于"我"的随机值
• 单次操作时间复杂度期望 \(O(logn)\)

怎么实现?

一些宏定义,方便写代码:

#define lc (tr[k].l)
#define rc (tr[k].r)
#define lson (tr[tr[k].l])
#define rson (tr[tr[k].r])

数组的定义:

• \(l\) , \(r\) -> 左右孩子
• \(v\) 权值
• \(sz\) 子树大小,包括节点本身
• \(rd\) 一个随机生成的值
• \(cnt\) 与当前节点权值相同的节点个数

struct TreapData{
    int l,r,v,sz,rd,cnt;
};

• \(update\),更新一个节点的 sz 信息:

void upd(int k){tr[k].sz=lson.sz+rson.sz+tr[k].cnt;}

• \(rotate\),旋转,分为左旋和右旋,下图是左旋,即把左孩子旋转上来,当前节点就是左孩子的右孩子

    // 左旋
    void lturn(int &k){
        int t=rc; rc=tr[t].l,tr[t].l=k;
        tr[t].sz=tr[k].sz;upd(k);k=t;
    }
    // 右旋
    void rturn(int &k){
        __R int t=lc; lc=tr[t].r,tr[t].r=k;
        tr[t].sz=tr[k].sz;upd(k);k=t;
    }

• \(insert\),插入一个元素,利用左边比"我"小,右边比"我"大,很容易就能够定位该元素的位置,插入后调整,使得 treap 保持堆的性质

    void ins(int &k,int x){
        if (!k){tr[k=++sz].sz=tr[k].cnt=1,tr[k].v=x,tr[k].rd=rand(),tr[k].l=tr[k].r=0;return;}
        tr[k].sz++;
        if (x==tr[k].v) tr[k].cnt++; else
        if (x<tr[k].v){ins(lc,x);if (lson.rd<tr[k].rd) rturn(k);} else
            {ins(rc,x);if (rson.rd<tr[k].rd) lturn(k);}
    }

• \(delete\),删除一个元素,将该元素一直旋转直到成为叶子节点后直接删除

    void del(int &k,int x){
        if (!k) return;
        if (x==tr[k].v){
            if (tr[k].cnt>1) tr[k].sz--,tr[k].cnt--; else
            if (!lc||!rc) k=lc+rc; else
            if (lson.rd<rson.rd) rturn(k),tr[k].sz--,del(rc,x); else
                lturn(k),tr[k].sz--,del(lc,x);
        } else tr[k].sz--,del(x<tr[k].v?lc:rc,x);
    }

• \(rank\),查找该元素的排名,即比当前数小的数的个数+1,就是查找一个节点,求其左子树大小+1

    int rnk(int k,int x){
        if (!k) return 0;
        if (x==tr[k].v) return lson.sz+1; else
        if (x<tr[k].v) return rnk(lc,x); else
            return lson.sz+tr[k].cnt+rnk(rc,x);
    }

• \(kth\),查找第 k 大的数,分 3 种情况考虑,第 k 大在左子树（右子树）还是在当前的节点

    int kth(int k,int x){
        if (!k) return 0;
        if (x<=lson.sz) return kth(lc,x); else
        if (x>lson.sz+tr[k].cnt) return kth(rc,x-lson.sz-tr[k].cnt); else
            return tr[k].v;
    }

• \(pre\),求当前数的前驱,向左走直到小于当前数,然后一直向右直到叶子节点

    void pre(int k,int x){
        if (!k) return;
        if (tr[k].v<x) ans=k,pre(rc,x); else
            pre(lc,x);
    }

• \(next\),求当前数的猴子向右走直到大于当前数,然后一直向左直到叶子节点

    void nxt(int k,int x){
        if (!k) return;
        if (x<tr[k].v) ans=k,nxt(lc,x); else
            nxt(rc,x);
    }

// 洛谷 3369
#include <cstdio>
#include <cstdlib>
#define MAXN 500001
#define __R register
#define rep(i,a,b) for(__R int i=(a);i<=(b);i++)
#define per(i,a,b) for(__R int i=(a);i>=(b);i--)
#define Rep(i,a,b) for(__R int i=(a);i<(b);i++)
#define Per(i,a,b) for(__R int i=(a);i>(b);i--)

int n,ans;

struct TreapNode{int l,r,v,rd,sz,cnt;};
struct TreapData{
    int rt,sz;
    TreapNode tr[MAXN];
#define lc (tr[k].l)
#define rc (tr[k].r)
#define lson (tr[tr[k].l])
#define rson (tr[tr[k].r])
    void upd(int k){tr[k].sz=lson.sz+rson.sz+tr[k].cnt;}
    void lturn(int &k){
        __R int t=rc; rc=tr[t].l,tr[t].l=k;
        tr[t].sz=tr[k].sz;upd(k);k=t;
    }
    void rturn(int &k){
        __R int t=lc; lc=tr[t].r,tr[t].r=k;
        tr[t].sz=tr[k].sz;upd(k);k=t;
    }
    void ins(int &k,int x){
        if (!k){tr[k=++sz].sz=tr[k].cnt=1,tr[k].v=x,tr[k].rd=rand(),tr[k].l=tr[k].r=0;return;}
        tr[k].sz++;
        if (x==tr[k].v) tr[k].cnt++; else
        if (x<tr[k].v){ins(lc,x);if (lson.rd<tr[k].rd) rturn(k);} else
            {ins(rc,x);if (rson.rd<tr[k].rd) lturn(k);}
    }
    void del(int &k,int x){
        if (!k) return;
        if (x==tr[k].v){
            if (tr[k].cnt>1) tr[k].sz--,tr[k].cnt--; else
            if (!lc||!rc) k=lc+rc; else
            if (lson.rd<rson.rd) rturn(k),tr[k].sz--,del(rc,x); else
                lturn(k),tr[k].sz--,del(lc,x);
        } else tr[k].sz--,del(x<tr[k].v?lc:rc,x);
    }
    int rnk(int k,int x){
        if (!k) return 0;
        if (x==tr[k].v) return lson.sz+1; else
        if (x<tr[k].v) return rnk(lc,x); else
            return lson.sz+tr[k].cnt+rnk(rc,x);
    }
    int kth(int k,int x){
        if (!k) return 0;
        if (x<=lson.sz) return kth(lc,x); else
        if (x>lson.sz+tr[k].cnt) return kth(rc,x-lson.sz-tr[k].cnt); else
            return tr[k].v;
    }
    void pre(int k,int x){
        if (!k) return;
        if (tr[k].v<x) ans=k,pre(rc,x); else
            pre(lc,x);
    }
    void nxt(int k,int x){
        if (!k) return;
        if (x<tr[k].v) ans=k,nxt(lc,x); else
            nxt(rc,x);
    }
} S;

int main(){
    srand(23333333);
    scanf("%d",&n);
    int opt,x;
    while (n--){
        scanf("%d%d",&opt,&x);
        switch (opt){
            case 1 : S.ins(S.rt,x); break;
            case 2 : S.del(S.rt,x); break;
            case 3 : printf("%d\n",S.rnk(S.rt,x)); break;
            case 4 : printf("%d\n",S.kth(S.rt,x)); break;
            case 5 : ans=0,S.pre(S.rt,x),printf("%d\n",S.tr[ans].v); break;
            case 6 : ans=0,S.nxt(S.rt,x),printf("%d\n",S.tr[ans].v); break;
        }
    }
    return 0;
}