[模板] 平衡树之Treap

正在大力统一码风, 先放代码

成功压行30行 可读性没有丝毫改变

// Treap
// 注意我们每一个会影响 Treap 结构的操作完成后都需要更新
// 这份代码的左子树小于右子树
// 值得注意的是在(关于左右子树大小关系与此程序相同) BST 上一个节点的左儿子一定小于当前位置, 但它并不一定是
// 小于当前位置的最大值, 这种情况只出现在其左儿子没有右儿子的情况下

# include <iostream>
# include <cstdio>
# include <queue>
# include <cstdlib>
# define MAXN 100005
# define INF (int)1<<25
using namespace std;

//
template<typename T>
T rd(){
	T fl = 1, x = 0;
	int ch = getchar();
	for(	; !isdigit(ch); ch = getchar()){
		if(ch == '-'){
			fl = -1;
		}
	}
	for(	;  isdigit(ch); ch = getchar()){
		x = x * 10 + ch - '0';
	}
	return x*fl;
}

//
struct Treap{
	int l, r;
	int val, cnt, siz; // 关键码(记录的值), 副本数, 子树大小
	// 这个 siz 与普通树的 siz 不同, 它是要计算上树上的每个点的 cnt 的
	int trp; // 当前节点 Treap 随机值
}t[MAXN]; 
int cntN; // 记录平衡树上的总节点数
int root; // 记录当前的树根

//
int NewN(int val){
	t[++cntN].val = val;
	t[cntN].trp = rand();
	t[cntN].cnt = t[cntN].siz = 1;
	return cntN;
}

void PushUp(int now){
	t[now].siz = t[t[now].l].siz + t[t[now].r].siz + t[now].cnt;
} // 操作后更新节点大小

void Build(){
	NewN(-INF), NewN(INF);
	root = 1, t[1].r = 2;
	PushUp(root);
} // 初始化先建立一个 +INF 节点和一个 -INF 节点

//
void LeftRot(int &now){
	int son = t[now].r;
	t[now].r = t[son].l, t[son].l = now, now = son;
	PushUp(t[now].l), PushUp(now);
} // 注意传入的是地址, 因为操作后原先的 x 点已经不在原来的位置了

void RightRot(int &now){
	int son = t[now].l;
	t[now].l = t[son].r, t[son].r = now, now = son;
	PushUp(t[now].r), PushUp(now);
} // 注意事项同上


void Insert(int &now, int val){
	if(!now){
		now = NewN(val);
		return;
	} // 这是当前树上无该节点的情况

	if(val == t[now].val){
		t[now].cnt += 1;
		PushUp(now);
		return;
	} // 在树上查找到该节点

	if(val < t[now].val){
		Insert(t[now].l, val);
		if(t[now].trp < t[t[now].l].trp){
			RightRot(now); // 注意是右旋	
		}
	} // 新增值小于当前值
	else{
		Insert(t[now].r, val);
		if(t[now].trp < t[t[now].r].trp){
			LeftRot(now); // 注意是左旋	
		}
	} // 新增值大于当前值
	
	PushUp(now); // 不要忘记更新
} // 同样需要传入地址, 因为可能涉及到旋转操作

void Delete(int &now, int val){
	if(!now){
		return;
	} // 不存在这个节点

	if(val == t[now].val){
		if(t[now].cnt > 1){
			t[now].cnt -= 1;
			PushUp(now);
			return;
		} // 存在副本的值只需减少一个副本
		else if(t[now].l || t[now].r){
			if((!t[now].r) || t[t[now].l].trp > t[t[now].r].trp){
				RightRot(now);
				Delete(t[now].r, val); // 注意这下我们查找到的位置被旋转到了右儿子位置
			} // 无右儿子的节点或是左旋后不满足大根堆性质的节点我们进行右旋
			else{
				LeftRot(now);
				Delete(t[now].l, val); // 这下我们查找的值被旋转到了左儿子
			}
			PushUp(now); // 对 Treap 的结构进行更改后不要忘记更新
		} // 非叶子节点需要旋转
		else{
			now = 0;
			return;
		} // 叶子节点直接删除
	} // 查找到了当前的值
	
	if(val < t[now].val){
		Delete(t[now].l, val);
	} // 没有找到, 继续向下查找
	else{
		Delete(t[now].r, val);
	}

	PushUp(now);
} // 同样需要传入地址

//
int AskRk(int now, int val){
	if(!now){
		return 0;
	} // 没有查找到这个值

	if(val == t[now].val){
		return t[t[now].l].siz + 1;
	} // 查到当前值, 注意因为每个值可能有多个副本, 所以我们要通过左儿子获取排名
	
	if(val < t[now].val){
		return AskRk(t[now].l, val);
	} // 没有查找到, 继续向下查找
	else{
		return AskRk(t[now].r, val) + t[t[now].l].siz + t[now].cnt;
	} // 注意这里要将前面的值的个数加上才是最终排名
} // 注意询问部分就不再需要传入地址了, 因为 Treap 中询问不会对节点进行操作
// 但是, 在 Splay 中, 询问会操作节点, 也就需要传地址了

int AskNum(int now, int rank){
	if(!now){
		return INF;
	}
	
	if(t[t[now].l].siz >= rank){
		return AskNum(t[now].l, rank);	
	} // 答案于左子树中

	if(t[t[now].l].siz + t[now].cnt >= rank){
		return t[now].val;
	} // 当前节点恰好是答案(此时已经排除了答案在左子树中的情况)

	else{
		return AskNum(t[now].r, rank-t[t[now].l].siz-t[now].cnt);
	} // 答案于右子树中
	// 不要忘记右子树的情况需要对 rank 操作一下, 有点类似与权值线段树
}

int AskPre(int now, int val){
	if(!now){
		return -INF;
	}
	else if(val <= t[now].val){
		return AskPre(t[now].l, val);
	}
	else{
		return max(t[now].val, AskPre(t[now].r, val));
	}
}

int AskNxt(int now, int val){
	if(!now){
		return INF;	
	}
	else if(val >= t[now].val){
		return AskNxt(t[now].r, val);
	}
	else{
		return min(t[now].val, AskNxt(t[now].l, val));
	}
}

//
int main(){
	Build();

	int n;
	
	n = rd<int>();
	
	for(int opt, x; n; n--){
		opt = rd<int>(), x = rd<int>();
		if(opt == 1){
			Insert(root, x);
		}
		else if(opt == 2){
			Delete(root, x);
		}
		else if(opt == 3){
			printf("%d\n", AskRk(root, x) - 1); // 注意这里需要减 1, 因为我们查到的排名受建树时的 -INF 影响
		}
		else if(opt == 4){
			printf("%d\n", AskNum(root, x + 1)); // 这里加 1 也是相同原因
		}
		else if(opt == 5){
			printf("%d\n", AskPre(root, x));
		}
		else if(opt == 6){
			printf("%d\n", AskNxt(root, x));
		}
	}

    return 0;
}