POJ 1442 Black Box ( Treap )

#include<iostream>
#include<cstdio>
#include<cstring>
#include<string>
#include<algorithm>
#include<cmath>
#include<queue>
#include<map>
#include<set>
#include<cstdlib>
using namespace std;

#define inf 0x3f3f3f3f
#define eps 1e-8
#define LL long long
#define ULL unsigned long long
#define MP make_pair
#define mxn 300030

struct node {
	node *ch[2];
	int v, r, s;
	node( int v ): v( v ) {
		r = rand();
		ch[0] = ch[1] = NULL;
		s = 1;
	}
	int cmp( int x ) {
		if( v == x )
			return -1;
		return x < v ? 0 : 1;
	}
	void mt() {
		s = 1;
		if( ch[0] )
			s += ch[0] -> s;
		if( ch[1] )
			s += ch[1] -> s;
	}
};
void rt( node *&p, int d ) {
	node *k = p -> ch[d^1];
	p -> ch[d^1] = k -> ch[d];
	k -> ch[d] = p;
	p -> mt();
	k -> mt();
	p = k;
}
void insert( node *&p, int x ) {
	if( p == NULL ) {
		p = new node( x );
		return;
	}
	int d = x < p -> v ? 0 : 1;
	insert( p -> ch[d], x );
	if( p -> ch[d] -> r > p -> r )
		rt( p, d ^ 1 );
	p -> mt();
}
bool find( node *p, int x ) {
	while( p ) {
		int d = p -> cmp( x );
		if( d == -1 )
			return 1;
		p = p -> ch[d];
	}
	return 0;
}
void remove( node *&p, int x ) {
	int d = p -> cmp( x );
	if( d == -1 ) {
		if( p -> ch[1] && p -> ch[0] ) {
			int d2 = p -> ch[0] -> r > p -> ch[1] -> r ? 1 : 0;
			rt( p, d2 );
			remove( p -> ch[d2], x );
		}
		else {
			node *u = p;
			if( p -> ch[0] == NULL )
				p = p -> ch[1];
			else
				p = p -> ch[0];
			delete u;
		}
	}
    else
            remove( p -> ch[d], x );
	if( p )
		p -> mt();
}
int kth( node *p, int k ) {
	if( p == NULL || k <= 0 || k > p -> s )
		return 0;
	int s = p -> ch[0] == NULL ? 0 : p -> ch[0] -> s;
	if( k == s + 1 )
		return p -> v;
	if( k <= s )
		return kth( p -> ch[0], k );
	return kth( p -> ch[1], k - s - 1 );
}
int rk( node *p, int x ) {
	int R;
	if( p -> ch[0] == NULL )
		R = 0;
	else
		R = p -> ch[0] -> s;
	if( x == p -> v )
		return R + 1;
	if( x < p -> v )
		return rk( p -> ch[0], x );
	return R + 1 + rk( p -> ch[1], x );
}
int a[mxn], cnt[mxn];
int n, m;
int main() {
//	freopen( "test.txt", "r", stdin );
	while( scanf( "%d%d", &n, &m ) != EOF ) {
		memset( cnt, 0, sizeof( cnt ) );
		for( int i = 1; i <= n; ++i ) {
			scanf( "%d", &a[i] );
		}
		for( int i = 1; i <= m; ++i ) {
			int t;
			scanf( "%d", &t );
			cnt[t]++;
		}
		node *root = NULL;
		int kk = 0;
		for( int i = 1; i <= n; ++i ) {
			insert( root, a[i] );
			while( cnt[i] ) {
				cnt[i]--;
				kk++;
				printf( "%d\n", kth( root, kk ) );
			}
		}
	}
	return 0;
}