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;
}