洛谷 P3369 【模板】普通平衡树

传送门

---

#include <iostream>
#include <cstring>
#include <cstdio>
#include <algorithm>
using namespace std;
//Mystery_Sky
//平衡树模板 
#define M 1000100
#define INF 0x3f3f3f3f
#define ll long long
inline int read()
{
	int x=0, f=1; char c=getchar();
	while(c<'0'||c>'9'){if(c=='-')f=-1;c=getchar();}
	while(c>='0'&&c<='9'){x=x*10+c-'0';c=getchar();}
	return x*f;
}
int n;
int root, tot;
int son[M][2];
int val[M], dat[M], size[M], cnt[M]; 

inline int New(int x)
{
	val[++tot] = x;
	dat[tot] = rand();
	size[tot] = 1;
	cnt[tot] = 1;
	return tot;
}

inline void pushup(int id)
{
	size[id] = size[son[id][0]] + size[son[id][1]] + cnt[id];
}

inline void build()
{
	root = New(-INF);
	son[root][1] = New(INF);
	pushup(root);
}

void Rotate(int &id, int d)//旋转操作
{
	int k = son[id][d^1];
	son[id][d^1] = son[k][d];
	son[k][d] = id;
	id = k;
	pushup(son[id][d]);
	pushup(id);
}

void insert(int &id, int x)//插入操作
{
	if(!id) {
		id = New(x);
		return;
	}
	if(x == val[id]) cnt[id]++;
	else {
		int d = x < val[id] ? 0 : 1;
		insert(son[id][d], x);
		if(dat[id] < dat[son[id][d]]) Rotate(id, d^1);
	}
	pushup(id);
}

void Remove(int &id, int x)//删除操作
{
	if(!id) return;
	if(x == val[id]) {
		if(cnt[id] > 1) {
			cnt[id]--;
			pushup(id);
			return;
		}
		if(son[id][0] || son[id][1]) {
			if(!son[id][1] || dat[son[id][0]] > dat[son[id][1]]) {
				Rotate(id, 1);
				Remove(son[id][1], x);
			}
			else Rotate(id, 0), Remove(son[id][0], x);
			pushup(id);
		}
		else id = 0;
		return;
	}
	x < val[id] ? Remove(son[id][0], x) : Remove(son[id][1], x);
	pushup(id);
	return;
}

int query_rank(int id, int x)//查询排名
{
	if(!id) return 0;
	if(x == val[id]) return size[son[id][0]] + 1;
	else if(x < val[id]) return query_rank(son[id][0], x);
	else return size[son[id][0]] + cnt[id] + query_rank(son[id][1], x);
}

int query_val(int id, int rank)//查询排名为rank对应的值
{
	if(!id) return INF;
	if(rank <= size[son[id][0]]) return query_val(son[id][0], rank);
	else if(rank <= size[son[id][0]] + cnt[id]) return val[id];
	else return query_val(son[id][1], rank - size[son[id][0]] - cnt[id]);
}

int query_pre(int x)//查询x的前驱
{
	int id = root;
	int ans;
	while(id) {
		if(val[id] < x) ans = val[id], id = son[id][1];
		else id = son[id][0];
	}
	return ans;
}

int query_next(int x)//查询x的后继
{
	int id = root;
	int ans;
	while(id) {
		if(val[id] > x) ans = val[id], id = son[id][0];
		else id = son[id][1];
	}
	return ans;
}

int main() {
	n = read();
	build();
	for(int i = 1; i <= n; i++) {
		int opt = read(), x = read();
		if(opt == 1) insert(root, x);
		else if(opt == 2) Remove(root, x);
		else if(opt == 3) printf("%d\n", query_rank(root, x) - 1);
		else if(opt == 4) printf("%d\n", query_val(root, x+1));
		else if(opt == 5) printf("%d\n", query_pre(x));
		else if(opt == 6) printf("%d\n", query_next(x));
	}
	return 0;
}