POJ 1442 Black Box [Treap]

Description

给两个数组
A
U
A表示插入到集合的数，U表示在插入第U[i]的数后询问
询问第K大的数
K每次询问之后加一

Algorithm

Treap实现的名次树的模板题
看算法竞赛入门经典训练指南

Code

#include <iostream>
#include <cstdlib>
#include <cstdio>
using namespace std;
struct Node {
        Node *ch[2]; //左右子树 0左1右
        int r; //优先级
        int v; //值
        int s; //子树
        Node(int v) : v(v) {
                ch[0] = ch[1] = NULL;
                r = rand();
                s = 1;
        }
        int cmp(int x) {
                if (x == v) return -1;
                return x < v ? 0 : 1;
        }
        void maintain() {
                s = 1;
                if (ch[0] != NULL) s += ch[0]->s;
                if (ch[1] != NULL) s += ch[1]->s;
        }
};

void rotate(Node* &o, bool d)
{
        Node *k = o->ch[!d];
        o->ch[!d] = k->ch[d];
        k->ch[d] = o;
        o->maintain();
        k->maintain();
        o = k;
}
void insert(Node *&o, int x) //在 o 节点插入值为 x的节点
{
        if (o == NULL) {
                o = new Node(x);
        } else {
                bool d = x > (o->v);
                insert(o->ch[d], x);
                if (o->ch[d]->r > o->r) rotate(o, !d);
        }
        o->maintain();
}
void remove(Node * &o, int x)
{
        int d = o->cmp(x);
        if (d == -1) {
                Node *u = o;
                if (o -> ch[0] != NULL && o->ch[1] != NULL) {
                        bool d2 = o->ch[0]->r < o->ch[1]->r;
                        rotate(o, d2);
                        remove(o->ch[d2], x);
                } else {
                        if (o->ch[0] == NULL) {
                                o = o->ch[1];
                        } else {
                                o = o->ch[0];
                                delete u;
                        }
                }
        } else {
                remove(o->ch[d], x);
                if (o != NULL) {
                        o->maintain();
                }
        }
}
int kth(Node *o, int k)
{
        if (o == NULL || k <= 0 || k > o->s) return 0;
        int s = o->ch[0] == NULL ? 0 : o->ch[0]->s;
        if (k == s + 1) {
                return o->v;
        } else {
                if(k <= s) {
                        return kth(o->ch[0], k);
                } else {
                        return kth(o->ch[1], k-s-1);
                }
        }
}
int n, m;
const int MAXM = 30000 + 9;
const int MAXN = MAXM;
int a[MAXM];
int u[MAXN];
void solve()
{
        for (int i = 0; i < m; i++) {
                scanf("%d", &a[i]);
        }
        for (int i = 0; i < n; i++) {
                scanf("%d", &u[i]);
        }
        int j = 0;
        Node *treap = NULL;
        for (int i = 0; i < n; i++) {
                for (; j < u[i]; j++) {
                        insert(treap, a[j]);
                }
                printf("%d\n", kth(treap, i + 1));
        }
}
int main()
{
        //freopen("in.txt", "r", stdin);
        while (scanf("%d%d", &m, &n) != EOF) {
                solve();
        }
}