笔记 数据结构 平衡树

#include <iostream>
#include <cstring>
#include <cstdio>
#include <algorithm>
using namespace std;
const int N = 100010, inf = 1e8;
int n;

struct node {
	int l, r;//左右儿子在tr上的坐标
	int key;//二叉搜索树的权值
	int val;//堆的权值，让搜索树变得随机，防止退化成一条链
	int cnt;//相同的数有多少个
	int size;//所有子树中有多少个数
} tr[N];
int root;//源点，搜索从这里开始
int idx;//记录当前最末端的节点

void pushup(int p) {//向上传递信息
	tr[p].size = tr[tr[p].l].size + tr[tr[p].r].size + tr[p].cnt;
}

int get_node(int key) { //创建节点
	tr[++idx].key = key;
	tr[idx].val = rand();
	tr[idx].cnt = 1;
	tr[idx].size = 1;
	return idx;
}

void build() {//建树
	get_node(-inf), get_node(inf); //插入两个哨兵,防止越界
	root = 1, tr[1].r = 2;
	pushup(root);
}
//p:x q:y 
void zig(int &p) { //右旋
	int q = tr[p].l;//
	tr[p].l = tr[q].r, tr[q].r = p;
	p = q;
	pushup(tr[p].r), pushup(p);
}

void zag(int &p) { //左旋
	int q = tr[p].r;
	tr[p].r = tr[q].l;
	tr[q].l = p;
	p = q;
	pushup(tr[p].l), pushup(p);
}

void insert(int &p, int key) {
	if (!p)
		p = get_node(key);//原本没有这个数，直接生成新节点即可
	else if (tr[p].key == key)//已经有了这个节点，cnt++
		tr[p].cnt++;
	else if (tr[p].key > key) {
		insert(tr[p].l, key);
		if (tr[tr[p].l].val > tr[p].val)//防止退化
			zig(p);
	} else if (tr[p].key < key) {
		insert(tr[p].r, key);
		if (tr[tr[p].r].val > tr[p].val)
			zag(p);
	}
	pushup(p);
}

void remove(int &p, int key) {
	if (!p)
		return ;
	if (tr[p].key == key) {
		if (tr[p].cnt > 1)
			tr[p].cnt--;
		else if (tr[p].l || tr[p].r) {//左右节点都存在
			if (!tr[p].r || tr[tr[p].l].val > tr[tr[p].r].val) {
				zig(p);//右旋，把这个点旋到叶节点
				remove(tr[p].r, key);
			} else {
				zag(p);//左旋，把这个点旋到叶节点
				remove(tr[p].l, key);
			}
		} else
			p = 0;//旋转到了叶节点,将p清除
	} else if (tr[p].key > key)
		remove(tr[p].l, key);
	else
		remove(tr[p].r, key);
	pushup(p);

}

int get_rank_by_key(int p, int key) { //通过数值找排名
	if (!p)
		return 0;
	if (tr[p].key == key)
		return tr[tr[p].l].size + 1;
	if (tr[p].key > key)
		return get_rank_by_key(tr[p].l, key); //在左边
	return tr[tr[p].l].size + tr[p].cnt + get_rank_by_key(tr[p].r, key); //在右边
}

int get_key_by_rank(int p, int rank) { //通过排名找数值
	if (!p)
		return inf;
	if (tr[tr[p].l].size >= rank)
		return get_key_by_rank(tr[p].l, rank);
	if (tr[tr[p].l].size + tr[p].cnt >= rank)
		return tr[p].key;//算上了当前数
	return get_key_by_rank(tr[p].r, rank - tr[tr[p].l].size - tr[p].cnt);
}

int get_prev(int p, int key) { //前驱,严格小于k的最大数
	if (!p)
		return -inf;
	if (tr[p].key >= key)
		return get_prev(tr[p].l, key);
	return max(tr[p].key, get_prev(tr[p].r, key)); //有可能是当前数
}

int get_next(int p, int key) { //后驱
	if (!p)
		return inf;
	if (tr[p].key <= key)
		return get_next(tr[p].r, key);
	return min(tr[p].key, get_next(tr[p].l, key));
}

int main() {
	build();
	cin >> n;
	while (n--) {
		int opt, x;
		cin >> opt >> x;
		if (opt == 1)
			insert(root, x);
		else if (opt == 2)
			remove(root, x); //有头部哨兵的存在，所以排名要变化
		else if (opt == 3)
			cout << get_rank_by_key(root, x) - 1 << endl;
		else if (opt == 4)
			cout << get_key_by_rank(root, x + 1) << endl;
		else if (opt == 5)
			cout << get_prev(root, x) << endl;
		else
			cout << get_next(root, x) << endl;
	}
}