Solution - P3369 【模板】普通平衡树

思路

用的 FHQ Treap。真的好用。

代码

#include <bits/stdc++.h>
#define rint register int
#define rllong register long long
#define llong long long
#define N 200005
using namespace std;

int val[N], pri[N], cnt[N], siz[N];
int ls[N], rs[N], top, root;
int q;

#define pushup(x) (siz[x] = siz[ls[x]]+siz[rs[x]]+cnt[x])

inline int make_node(rint x){
	++top;
	val[top] = x, pri[top] = rand();
	cnt[top] = siz[top] = 1;
	return top;
}

inline void split(rint x, rint& L, rint& R, rint k){
	if(x == 0){
		L = R = 0;
		return;
	}
	if(val[x] <= k){
		L = x;
		split(rs[x], rs[L], R, k);
	}
	else{
		R = x;
		split(ls[x], L, ls[R], k);
	}
	pushup(L), pushup(R);
	return;
}

inline int merge(rint L, rint R){
	if(!L || !R) return L|R;
	rint res;
	if(pri[L] > pri[R]){
		rs[L] = merge(rs[L], R);
		res = L;
	}
	else{
		ls[R] = merge(L, ls[R]);
		res = R;
	}
	pushup(res);
	return res;
}

inline void insert(rint k){
	rint L, M, R;
	split(root, L, R, k);
	split(L, L, M, k-1);
	if(M) ++cnt[M], ++siz[M];
	else M = make_node(k);
	L = merge(L, M);
	root = merge(L, R);
	return;
}

inline void remove(rint k){
	rint L, M, R;
	split(root, L, R, k);
	split(L, L, M, k-1);
	if(cnt[M] > 1) --cnt[M], --siz[M];
	else M = 0;
	L = merge(L, M);
	root = merge(L, R);
	return;
}

inline int query1(rint k){ // Return rank of a number
	rint L, R, res;
	split(root, L, R, k-1);
	res = siz[L] + 1;
	root = merge(L, R);
	return res;
}

inline int query2(rint k){
    rint x = root, rnk = 0;
    while (x){
        if(k <= rnk+siz[ls[x]]) 
			x = ls[x];
		else if(k > rnk+siz[ls[x]]+cnt[x])
            rnk += siz[ls[x]]+cnt[x], x = rs[x];
        else break;
    }
    return val[x];
}

inline int query3(rint k){ // Get maxinum in prefix
	rint L, R, x, res;
	split(root, L, R, k-1);
	x = L;
	while(rs[x]) x = rs[x];
	res = val[x];
	root = merge(L, R);
	return res;
}

inline int query4(rint k){ // Get mininum in suffix
	rint L, R, x, res;
	split(root, L, R, k);
	x = R;
	while(ls[x]) x = ls[x];
	res = val[x];
	root = merge(L, R);
	return res;
}

int main(){
	srand(time(0)^clock());
	scanf("%d", &q);
	while(q--){
		rint op, k;
		scanf("%d %d", &op, &k);
		switch(op){
			case 1: insert(k); break;
			case 2: remove(k); break;
			case 3: printf("%d\n", query1(k)); break;
			case 4: printf("%d\n", query2(k)); break;
			case 5: printf("%d\n", query3(k)); break;
			case 6: printf("%d\n", query4(k)); break;
		}
	}
	return 0;
}