P1801 黑匣子（Treap树）

待修改的数组查询第k大，简单题

#include<bits/stdc++.h>
using namespace std;
const int maxn=2e5+100;
int a[maxn];
int u[maxn];
int n,m;

struct Treap_tree {
    int ch[2];
    int v;
    int dat;
    int size;
    int cnt;
}t[maxn];
int tot;
int root;
int newNode (int v) {
    tot++;
    t[tot].v=v;
    t[tot].dat=rand();//随机优先级
    t[tot].size=1;
    t[tot].cnt=1;
    return tot; 
} 
void pushup (int x) {
    t[x].size=t[t[x].ch[0]].size+t[t[x].ch[1]].size+t[x].cnt;
}

void rotate (int &id,int d) {
    int tt=t[id].ch[d^1];
    t[id].ch[d^1]=t[tt].ch[d];
    t[tt].ch[d]=id;
    id=tt;
    pushup(t[id].ch[d]);
    pushup(id);
}
void ins (int &id,int v) {
    if (!id) {
        id=newNode(v);
        return;
    }
    if (v==t[id].v) t[id].cnt++;
    else {
        ins(t[id].ch[v>t[id].v],v);
        if (t[id].dat<t[t[id].ch[v>t[id].v]].dat) rotate(id,v<t[id].v);
    }
    pushup(id);
}
int kth (int id,int k) {
    if (!id) return 1e9;
    if (k<=t[t[id].ch[0]].size)
        return kth(t[id].ch[0],k);
    else if (k<=t[t[id].ch[0]].size+t[id].cnt)
        return t[id].v;
    else
        return kth(t[id].ch[1],k-t[t[id].ch[0]].size-t[id].cnt);
}
int main () {
    scanf("%d%d",&n,&m);
    for (int i=1;i<=n;i++) scanf("%d",a+i);
    for (int i=0;i<m;i++) scanf("%d",u+i);
    int tot=0;
    for (int i=1;i<=n;i++) {
        ins(root,a[i]);
        while (i==u[tot]) {
            tot++;
            printf("%d
",kth(root,tot));
        }
    }
}