BZOJ 3224 - 用Treap实现名次树

码了一发名次树，然后在remove和rank上GG了…… remove的话换了一种更保险的写法；而rank直接抄了Rujia Liu的代码…… 给Rj L跪了…

// BZOJ 3224, Treap + Kth

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cstdlib>
using namespace std;
 
 #define read(x) scanf("%d", &x)

 struct Node {
 	Node *son[2]; // 0 <=> left, 1 <=> right
 	int p, v, s;
 	int cmp(int x) {
 		if (v==x) return -1;
 		return x<v ? 0 : 1;
 	}
 	void maintain() { s = son[1]->s + son[0]->s +1; }
 } *root;

 Node *null = new Node();

 void init() {
 	null->s = 0;
 	root = null;
 }

 void rotate(Node *&o, int d) {
 	Node *k = o->son[d^1];
 	o->son[d^1] = k->son[d];
 	k->son[d] = o;
 	o->maintain();
 	k->maintain();
 	o = k;
 } 

 // Movement 1
 void insert(Node *&o, int x) {
 	if (o==null) {
 		o = new Node();
 		o->v = x;
 		o->son[0] = o->son[1] = null;
 		o->p = rand();
 	}
 	else {
 		int d = (x < o->v ? 0 : 1); // 也可以直接用cmp函数，但后面要写得麻烦一点
 		insert(o->son[d], x);
 		if (o->son[d]->p > o->p) rotate(o, d^1);
 	}
 	o->maintain();
 }

 // Movement 2
 void remove(Node *&o, int x) {
 	if (o==null) return;
 	int d = o->cmp(x);
 	if (d==-1) {
 		Node *u = o;
 		if (o->son[0] != null && o->son[1] != null) {
 			int d2 = (o->son[0]->p > o->son[1]->p ? 1 : 0);
 			rotate(o, d2);
 			remove(o->son[d2], x);
 		}
 		else {
 			if (o->son[0] == null) o = o->son[1]; else o = o->son[0];
 			delete u;
 		}
 	} else remove(o->son[d], x);
 	if (o!=null) o->maintain();
 } 

 // Movement 3
 int get_rank(Node *o, int x) {
    if (o == null) return 1;
    if (x <= o->v) return get_rank(o->son[0], x);
    return get_rank(o->son[1], x) + (o->son[0] == null ? 0 : o->son[0]->s) + 1;
 }

 // Movement 4
 int kth(Node *o, int k) {
 	if (o==null || k<=0 || k > o->s) return 0; // 鲁棒语句：没有第k大的元素
 	int s = o->son[0]->s;
 	if (k==s+1) return o->v;
 	else if (k<=s) return kth(o->son[0], k);
 	else return kth(o->son[1], k-s-1);
 }

 // Movement 5
 int query_pre(int x) {  // 注意求前驱或后继时要忽略相同元素
 	Node *t = root;
 	int ret=0;
 	while (t != null) {
 		if (t->v < x) {
 			ret = t->v;
 			t = t->son[1];
 		} else t = t->son[0];
 	}
 	return ret;
 }

 // Movement 6
 int query_suf(int x) {
 	Node *t = root;
 	int ret=0;
 	while (t != null) {
 		if (t->v > x) {
 			ret = t->v;
 			t = t->son[0];
 		} else t = t->son[1];
 	}
 	return ret;
 }

int main()
{
	int m, op, x;
    read(m);
    init();
    while (m--) {
    	read(op); read(x);
    	if (op==1) insert(root, x);
    	else if (op==2) remove(root, x);
    	else if (op==3) printf("%d\n", get_rank(root, x));
    	else if (op==4) printf("%d\n", kth(root, x));
    	else if (op==5) printf("%d\n", query_pre(x));
    	else printf("%d\n", query_suf(x));
    }

	return 0;
}