平衡树

P3369 【模板】普通平衡树

https://www.luogu.org/problemnew/show/P3369

BST

struct BST {
    int l, r; // 左右子节点在数组中的下标
    int val;  // 节点关键码
} a[SIZE]; // 数组模拟链表
int tot, root, INF = 1<<30;

int New(int val) {
    a[++tot].val = val;
    return tot;
}

void Build() {
    New(-INF), New(INF);
    root = 1, a[1].r = 2;
}

int Get(int p, int val) {
    if (p == 0) return 0; // 检索失败
    if (val == a[p].val) return p; // 检索成功
    return val < a[p].val ? Get(a[p].l, val) : Get(a[p].r, val);
}

void Insert(int &p, int val) {
    if (p == 0) {
        p = New(val); // 注意p是引用，其父节点的l或r值会被同时更新
        return;
    }
    if (val == a[p].val) return;
    if (val < a[p].val) Insert(a[p].l, val);
    else Insert(a[p].r, 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 ans;
}

void Remove(int val) {
    // 检索val，得到节点p
    int &p = root;
    while (p) {
        if (val == a[p].val) break;
        p = val < a[p].val ? a[p].l : a[p].r;
    }
    if (p == 0) return;
    if (a[p].l == 0) { // 没有左子树
        p = a[p].r; // 右子树代替p的位置，注意p是引用
    }
    else if (a[p].r == 0) { // 没有右子树
        p = a[p].l; // 左子树代替p的位置，注意p是引用
    }
    else { // 既有左子树又有右子树
        // 求后继节点
        int next = a[p].r;
        while (a[next].l > 0) next = a[next].l;
        // next一定没有左子树，直接删除
        Remove(a[next].val);
        // 令节点next代替节点p的位置
        a[next].l = a[p].l, a[next].r = a[p].r;
        p = next; // 注意p是引用
    }
}

treap

利用随机权值维护堆的形态

const int N=100010,INF=1<<30;
struct treap{int l,r,c,dat,cnt,size;}a[N];
int m,tot,root;

int newnode(int val)
{
    a[++tot].c=val;
    a[tot].dat=rand();
    a[tot].size=a[tot].cnt=1;
    return tot;
}

void update(int x)
{
    a[x].size=a[a[x].l].size+a[a[x].r].size+a[x].cnt;
}

void build()
{
    newnode(-INF),newnode(INF);
    root=1,a[1].r=2;
    update(root);
}

void zlg(int &p)//zag
{
    int q=a[p].r;
    a[p].r=a[q].l,a[q].l=p,p=q;
    update(a[p].l),update(p);
}
void zrg(int &p)//zig
{
    int q=a[p].l;
    a[p].l=a[q].r,a[q].r=p,p=q;
    update(a[p].r),update(p);
}

void insert(int &p,int val)
{
    if(p==0) {p=newnode(val);return;}
    if(a[p].c==val) {++a[p].cnt,++a[p].size;return;}
    if(val<a[p].c) 
    {
        insert(a[p].l,val);
        if(a[p].dat<a[a[p].l].dat) zrg(p);
    }
    else
    {
        insert(a[p].r,val);
        if(a[p].dat<a[a[p].r].dat) zlg(p);
    }
    update(p);
}

int grbv(int p,int val)
{
    if(p==0) return 0;
    if(a[p].c==val) return a[a[p].l].size+1;
    if(a[p].c>val) return grbv(a[p].l,val);
    return grbv(a[p].r,val)+a[a[p].l].size+a[p].cnt;
}

int gvbr(int p,int ran)
{
    if(p==0) return INF;
    if(a[a[p].l].size>=ran) return gvbr(a[p].l,ran);
    else if(a[a[p].l].size+a[p].cnt>=ran) return a[p].c;
    else return gvbr(a[p].r,ran-a[a[p].l].size-a[p].cnt);
}

int getnext(int val)
{
    int p=root,ans=2;
    while(p)
    {
        if(a[p].c==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].c>val&&a[p].c<a[ans].c) ans=p;
        if(val<a[p].c) p=a[p].l;
        else p=a[p].r;
    }
    return a[ans].c;
}

int getlast(int val)
{
    int p=root,ans=1;
    while(p)
    {
        if(a[p].c==val)//||(a[p].c>val&&a[p].l==0)
        {
            if(a[p].l>0)
            {
                p=a[p].l;
                while(a[p].r>0) p=a[p].r;
                ans=p;
            } 
            break;
        }
        if(a[p].c<val&&a[p].c>a[ans].c) ans=p;
        if(val<a[p].c) p=a[p].l;
        else p=a[p].r;
    }
    return a[ans].c;
}

void remove(int &p,int val)
{
    if(p==0) return;
    if(a[p].c==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)
                zrg(p),remove(a[p].r,val);
            else zlg(p),remove(a[p].l,val);
            update(p);
        }
        else p=0;
        return;
    }
    if(val<a[p].c) remove(a[p].l,val);
    else remove(a[p].r,val);
    update(p);
}

int main()
{
    srand(time(NULL));
    cin>>m;build();
    FOR(i,1,m)
    {
        int opt,x;cin>>opt>>x;
        switch(opt)
        {
            case 1:insert(root,x);break;
            case 3:cout<<grbv(root,x)-1<<endl;break;
            case 4:cout<<gvbr(root,x+1)<<endl;break;
            case 5:cout<<getlast(x)<<endl;break;
            case 6:cout<<getnext(x)<<endl;break;
            case 2:remove(root,x);break;
        }
    }
    return 0;
}