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;
}