萌新刚学 OI,20分求助Treap
查看原帖
萌新刚学 OI,20分求助Treap
254733
Night_Bringer楼主2021/3/26 21:48

rt

#include <cstdio>
#include <cstdlib>
#define INF 0x3f3f3f3f
const int MAXN = 1e6 + 5;
struct Treap {
	int Left_Child, Right_Child;
	int Val, Dat, Cnt, Siz;
	#define LC(x) tree[x].Left_Child
	#define RC(x) tree[x].Right_Child
	#define V(x) tree[x].Val
	#define D(x) tree[x].Dat
	#define C(x) tree[x].Cnt
	#define S(x) tree[x].Siz
};
Treap tree[MAXN];
int n;
int tot, root;
int New(int val) {
	tot++;
	V(tot) = val;
	D(tot) = rand();
	C(tot) = S(tot) = 1;
	return tot;
}
void Update(int pos) {
	S(pos) = S(LC(pos)) + S(RC(pos)) + C(pos);
}
void Build() {
	root = New(-INF);
	RC(1) = New(INF);
	Update(root);
}
int GetRankByVal(int pos, int val) {
	if(!pos)
		return -2;
	if(val == V(pos))
		return S(LC(pos)) + 1;
	if(val < V(pos))
		return GetRankByVal(LC(pos), val);
	return GetRankByVal(RC(pos), val) + S(LC(pos)) + C(pos);
}
int GetValByRank(int pos, int rank) {
	if(!pos)
		return INF;
	if(S(LC(pos)) >= rank)
		return GetValByRank(LC(pos), rank);
	if(S(LC(pos)) + C(pos) >= rank)
		return V(pos);
	return GetValByRank(RC(pos), rank - S(LC(pos)) - C(pos));
}
void Zig(int &pos) {
	int Lson = LC(pos);
	LC(pos) = RC(Lson);
	RC(Lson) = pos;
	pos = Lson;
	Update(RC(pos));
	Update(pos);
}
void Zag(int &pos) {
	int Rson = RC(pos);
	RC(pos) = LC(Rson);
	LC(Rson) = pos;
	pos = Rson;
	Update(LC(pos));
	Update(pos);
}
void Insert(int &pos, int val) {
	if(!pos) {
		pos = New(val);
		return;
	}
	if(val == V(pos)) {
		C(pos)++;
		Update(pos);
		return;
	}
	else if(val < V(pos)) {
		Insert(LC(pos), val);
		if(D(pos) < D(LC(pos)))
			Zig(pos);
	}
	else {
		Insert(RC(pos), val);
		if(D(pos) < D(RC(pos)))
			Zag(pos);
	}
	Update(pos);
}
int Get_Pre(int val) {
	int pos = root, res;
	while(pos) {
		if(V(pos) < val) {
			res = V(pos);
			pos = RC(pos);
		}
		else
			pos = LC(pos);
	}
	return res;
}
int Get_Next(int val) {
	int pos = root, res;
	while(pos) {
		if(V(pos) > val) {
			res = V(pos);
			pos = LC(pos);
		}
		else
			pos = RC(pos);
	}
	return res;
}
void Remove(int &pos, int val) {
	if(!pos)
		return;
	if(val == V(pos)) {
		if(C(pos) > 1) {
			C(pos)--;
			Update(pos);
			return;
		}
		if(LC(pos) || RC(pos)) {
			if(!RC(pos) || D(RC(pos)) > D(LC(pos))) {
				Zig(pos);
				Remove(pos, val);
			}
			else {
				Zag(pos);
				Remove(pos, val);
			}
		}
		else
			pos = 0;
		return;
	}
	else if(val < V(pos))
		Remove(LC(pos), val);
	else
		Remove(RC(pos), val);
	Update(pos);
}
int main() {
	Build();
	scanf("%d", &n);
	for(int i = 1, opt, a; i <= n; i++) {
		scanf("%d %d", &opt, &a);
		if(opt == 1)
			Insert(root, a);
		else if(opt == 2)
			Remove(root, a);
		else if(opt == 3)
			printf("%d\n", GetRankByVal(root, a) - 1);
		else if(opt == 4)
			printf("%d\n", GetValByRank(root, a + 1));
		else if(opt == 5)
			printf("%d\n", Get_Pre(a));
		else
			printf("%d\n", Get_Next(a));
	}
	return 0;
}
2021/3/26 21:48
加载中...