SPOJ 3273 Order statistic set

Treap 模板题

题意： 要求实现一种数据结构，支持：

Insert： 插入原来不存在的元素。

Remove： 删除原来存在的元素。

Kth： 询问第K大的元素，询问不合法(即k大于元素数量)，输出“invalid”。

Count：询问小于x的元素的数量。

解法： Treap的基本操作全齐了，红果果的裸题。

训练指南版

#include <stdio.h>
#include <string.h>
#include <algorithm>
#include <stdlib.h>
#include <time.h>
using namespace std;

// Treap, 大根堆写法,指针写法
struct Node {
	Node *ch[2];
	int r;
	int v;
	int size;
	int cmp(int x) const {
		if(x == v) return -1;
		return x < v ? 0 : 1;
	}
	void maintain() {
		size = ch[0]->size + ch[1]->size + 1;
	}
};

Node* nill;


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

void insert(Node* &o, int x) {
	if(o == nill) {
		o = new Node();
		o->ch[0] = o->ch[1] = nill;
		o->v = x;
		o->r = rand();
	}
	else {
		int d = o->cmp(x);
		insert(o->ch[d],x);
		if(o->ch[d]->r > o->r) rotate(o,d^1);
	}
	o->maintain();
}

void remove(Node* &o, int x) {
	int d = o->cmp(x);
	if(d == -1) {
		if(o->ch[0] == nill || o->ch[1] == nill) {
			Node* u = o;
			if(o->ch[0] == nill) o = o->ch[1];
			else  o = o->ch[0];
			delete u;
		}
		else {
			int d2 = (o->ch[0]->r > o->ch[1]->r ? 1 : 0);
			rotate(o,d2);
			remove(o->ch[d2],x);
		}
	}
	else remove(o->ch[d],x);
	if(o != nill) o->maintain();
}

int find(Node* o, int x) {
	while(o != nill) {
		int d = o->cmp(x);
		if(d == -1) return 1;
		else o = o->ch[d];
	}
	return 0;
}

int count(Node* o, int x) {
	int ret = 0;
	while(o != nill) {
		if(x > o->v) {
			ret += o->ch[0]->size + 1;
			o = o->ch[1];
		}
		else o = o->ch[0];
	}
	return ret;
}

int Kth(Node* o, int x) {
	while(o != nill) {
		if(o->ch[0]->size == x-1) return o->v;
		if(o->ch[0]->size >= x) o = o->ch[0];
		else {
			x -= o->ch[0]->size + 1;
			o = o->ch[1];
		}
	}
	return -1;
}

void init() {
	nill = new Node();
	nill->ch[0] = nill->ch[1] = nill;
	nill->size = 0;
	srand(time(NULL));
}

int main() {
	init();
	Node* root = nill;
	int q,x;
	char s[2];
	scanf("%d", &q);
	while(q--) {
		scanf("%s%d", s, &x);
		if(s[0] == 'I' && find(root,x) == 0) insert(root,x);
		else if(s[0] == 'D' && find(root,x) == 1) remove(root,x);
		else if(s[0] == 'C') printf("%d\n",count(root,x));
		else if(s[0] == 'K') {
			if(root->size < x) printf("invalid\n");
			else printf("%d\n", Kth(root,x));
		}
	}
	return 0;
}