普通平衡树

题目描述

思路

Treap 模板题

代码

#include <cstdio>
#include <cstdlib>
#include <ctime>

const int MAX = 110000 + 5;
int n, m, rt, tot, inf = 0x7f3f3f3f, ot;
char str[100];
struct Node {
	int lc, rc, key, size, cnt, pri;
	#define lc(x) t[x].lc
	#define rc(x) t[x].rc
	#define pri(x) t[x].pri
	#define key(x) t[x].key
	#define size(x) t[x].size
	#define cnt(x) t[x].cnt
} t[MAX];
inline void write(int x) {
	if (x == 0) { putchar('0'); return; }
	if (x < 0) putchar('-'), x = -x;
	ot = 0;
	while (x) str[++ot] = x % 10 + '0', x /= 10;
	while (ot) putchar(str[ot--]);
}
inline int read() {
	int s = 0, f = 1;
	char ch = getchar();
	while (ch < '0' || ch > '9') {
		if (ch == '-') f = -1;
		ch = getchar();
	}
	while (ch >= '0' && ch <= '9') s = s * 10 + ch - '0', ch = getchar();
	return s * f;
}
void show(int x) {
	printf("%d %d %d %d %d %d\n", lc(x), rc(x), pri(x), key(x), size(x), cnt(x));
	if (lc(x)) show(lc(x));
	if (rc(x)) show(rc(x));
}
void pushup(int r) { 
	size(r) = size(lc(r)) + size(rc(r)) + cnt(r); 
}
void zig(int &r) {
	int s = lc(r);
	lc(r) = rc(s);
	rc(s) = r;
	size(s) = size(r);
	pushup(r);
	r = s;
}
void zag(int &r) {
	int s = rc(r);
	rc(r) = lc(s);
	lc(s) = r;
	size(s) = size(r);
	pushup(r);
	r = s;
}
void insert(int &r, int k) {
	if (!r) {
		r = ++tot; pri(r) = rand(); key(r) = k;
		lc(r) = rc(r) = 0;
		size(r) = cnt(r) = 1;
		return;
	} else ++size(r);
	if (key(r) == k) ++cnt(r);
	else if (k < key(r)) {
		insert(lc(r), k);
		if (pri(lc(r)) < pri(r)) zig(r);
	} else {
		insert(rc(r), k);
		if (pri(rc(r)) < pri(r)) zag(r);
	}
}

void del(int &r, int k) {
	if (key(r) == k) {
		if (cnt(r) >= 2) --cnt(r), --size(r);
		else if (!lc(r) || !rc(r)) r = lc(r) + rc(r);
		else if (pri(lc(r)) < pri(rc(r))) {
			zig(r), del(r, k);
		} else {
			zag(r), del(r, k);
		}
		return;
	}
	--size(r);
	if (k < key(r)) del(lc(r), k);
	else del(rc(r), k);
}

int queryPre(int k) {
	int r = rt, res = inf;
	while (r) {
		if (k > key(r)) res = key(r), r = rc(r);
		else r = lc(r);
	}
	return res;
}

int queryNxt(int k) {
	int r = rt, res = inf;
	while (r) {
		if (k < key(r)) res = key(r), r = lc(r);
		else r = rc(r);
	}
	return res;
}
int queryKth(int k) {
	int r = rt, res = inf;
	while (r) {
		if (size(lc(r)) < k && size(lc(r)) + cnt(r) >= k) return key(r);
		else if (size(lc(r)) >= k) r = lc(r);
		else k -= size(lc(r)) + cnt(r), r = rc(r);
	}
	return res;
}
int queryRand(int k) {
	int r = rt, res = 0;
	while (r) {
		if (k == key(r)) return res + size(lc(r)) + 1;
		else if (k < key(r)) r = lc(r);
		else res += size(lc(r)) + cnt(r), r = rc(r);
	}
	return res;
}

int main() {
	srand(time(NULL));
	n = read();
	for (int i = 1, p, q; i <= n; ++i) {
		p = read(), q = read();
		// printf("%d %d\n", p, q);
		switch(p) {
			case 1: insert(rt, q); break;
			case 2: del(rt, q); break;
			case 3: m = queryRand(q); write(m); puts(""); break;
			case 4: m = queryKth(q); write(m); puts(""); break;
			case 5: m = queryPre(q); write(m); puts(""); break;
			case 6: m = queryNxt(q); write(m); puts("");
		}
		// show(rt);
	}
	return 0;
}