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

题链

Treap平衡树

zw:别把线段树宏定义的ls rs用在平衡树上... 没发现难找

#include <bits/stdc++.h>
//#include <ext/pb_ds/priority_queue.hpp>
//#pragma GCC optimize("O2")
using namespace std;
//using namespace __gnu_pbds;
#define Combine Pair, greater<Pair>, pairing_heap_tag
#define LL long long
#define ll long long
#define Pair pair<double,LL>
#define ULL unsigned long long
#define ls rt<<1
#define rs rt<<1|1
#define one first
#define two second
#define MS 1000009
#define INF 1e9
#define DBINF 1e100
#define Pi acos(-1.0)
#define eps 1e-9
#define mod 99999997
#define mod1 39989
#define mod2 1000000000

int n,m;
struct node{
	int l,r; // 左右孩子编号 
	int val,dat; // 键值，用于模仿堆的权值 
	int size,cnt;// 以该节点为根的树的大小，该节点的数的数量 
}p[MS];
int tot,root; // 总编号，根节点编号 

int add(int val){ // 新增节点返回节点编号 
	p[++tot] = {0,0,val,rand(),1,1};
	return tot;
}

void build(){ // 设置初始两个节点，正负无穷 
	root = add(-INF);
	p[root].r = add(INF);
}

void push_up(int rt){ // 更新 
	p[rt].size = p[p[rt].l].size + p[p[rt].r].size + p[rt].cnt;
}

int rank_val(int rt,int rank){ // 排名 -> 值 
	if(rt == 0) return INF;
	if(p[p[rt].l].size >= rank) return rank_val(p[rt].l,rank);
	if(p[p[rt].l].size + p[rt].cnt >= rank) return p[rt].val;
	return rank_val(p[rt].r,rank - p[p[rt].l].size - p[rt].cnt);
}

int val_rank(int rt,int val){ // 值 -> 排名 
	if(rt == 0) return 0;
	if(p[rt].val == val) return p[p[rt].l].size + 1;
	if(p[rt].val > val) return val_rank(p[rt].l,val);
	return p[p[rt].l].size + p[rt].cnt + val_rank(p[rt].r,val);
}

void left_rotate(int &rt){ // 左旋 逆时针 
	int rrs = p[rt].r;
	p[rt].r = p[rrs].l;
	p[rrs].l = rt;
	rt = rrs;
	push_up(p[rt].l);
	push_up(rt);
}

void right_rotate(int &rt){ // 右旋 顺时针 
	int lls = p[rt].l;
	p[rt].l = p[lls].r;
	p[lls].r = rt;
	rt = lls;
	push_up(p[rt].r);
	push_up(rt);
}

void insert(int &rt,int val){ // 插入 rt为引用 根节点的位置可能会变
	if(rt == 0){ // 无该值则新增 
		rt = add(val);
		return;
	}
	if(p[rt].val == val){ // 有该值则数量++ 
		p[rt].cnt++;
	}
	else if(val < p[rt].val){
		insert(p[rt].l,val);
		if(p[p[rt].l].dat > p[rt].dat) right_rotate(rt); // 右旋以满足大根堆性质 
	}
	else{
		insert(p[rt].r,val);
		if(p[p[rt].r].dat > p[rt].dat) left_rotate(rt); // 左旋以满足大根堆性质 
	}
	push_up(rt); // 不忘更新 
}

int getfront(int rt,int val){ //前驱 
	int ans = 1; // p[1].val  = -INF;
	while(rt){
		if(p[rt].val == val){ // 找到对应值 
			if(p[rt].l != 0){ // 寻找左子树最大值 
				rt = p[rt].l;
				while(p[rt].r != 0) rt = p[rt].r;
				ans = rt;
			}
			break;
		}
		if(p[ans].val < p[rt].val && p[rt].val < val) ans = rt; // 边查找边更新 
		if(val < p[rt].val) rt = p[rt].l;
		else rt = p[rt].r;
	}
	return p[ans].val;
}

int getnext(int rt,int val){ // 后继 
	int ans = 2;
	while(rt){
		if(p[rt].val == val){ // 找到对应值 
			if(p[rt].r != 0){ // 寻找右子树最小值 
				rt = p[rt].r;
				while(p[rt].l != 0) rt = p[rt].l;
				ans = rt;
			}
			break;
		}
		if(val < p[rt].val && p[rt].val < p[ans].val) ans = rt; // 边查找边更新 
		if(val < p[rt].val) rt = p[rt].l;
		else rt = p[rt].r;
	}
	return p[ans].val;
}

void remove(int &rt,int val){ //删除 rt为引用 根节点的位置可能会变 
	if(rt == 0) return;
	if(p[rt].val == val){
		if(p[rt].cnt > 1){ // 若有多个值则删数量即可 
			p[rt].cnt--;
			push_up(rt);
			return;
		}
		if(p[rt].l || p[rt].r){ // 使需要被删除的数向下旋 到达叶子节点即可直接删除 左右旋保持堆性质 
			if(p[rt].r == 0 || p[p[rt].l].dat > p[p[rt].r].dat){
				right_rotate(rt);
				remove(p[rt].r,val);
			}
			else{
				left_rotate(rt);
				remove(p[rt].l,val);
			}
			push_up(rt);
		}
		else rt = 0; // 直接删除叶子节点 
		return;
	}
	if(val < p[rt].val) remove(p[rt].l,val);
	else remove(p[rt].r,val);
	push_up(rt);
}

int main() {
	ios::sync_with_stdio(false);
	cin >> n;
	build();
	while(n--){
		int op,x;
		cin >> op >> x;
		if(op == 1) insert(root,x);
		else if(op == 2) remove(root,x);
		else if(op == 3) cout << val_rank(root,x)-1 << endl; // 初始时有一个无穷小 
		else if(op == 4) cout << rank_val(root,x+1) << endl; // 初始时有一个无穷小 
		else if(op == 5) cout << getfront(root,x) << endl;
		else cout << getnext(root,x) << endl;
	}
	
	


	return 0;
}