SP10628 COT - Count on a tree 主席树

Code:

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<string>
#include<iostream>
 
using namespace std;
 
void SetIO(string a){
	string in = a + ".in";
	freopen(in.c_str(),"r",stdin);
}
 
void debug(){
	cout << 233 << endl;
}
 
const int maxn = 100000 + 5;
 
int n, m;
 
int val[maxn];
 
int Sorted[maxn];
 
inline void Disperse(){
	sort(Sorted + 1, Sorted + 1 + n);
	for(int i = 1;i <= n; ++i)
		val[i] = lower_bound(Sorted + 1, Sorted + 1 + n, val[i]) - Sorted;
}
 
int head[maxn << 1], to[maxn << 1], nex[maxn << 1], edges;
 
inline void add_edge(int u, int v){
	nex[++edges] = head[u];
	head[u] = edges;
	to[edges] = v;
}
 
inline void Read(){
	scanf("%d%d",&n, &m);
	for(int i = 1;i <= n; ++i){
		scanf("%d",&val[i]), Sorted[i] = val[i];
	}
 
	for(int i = 1;i < n; ++i){
		int a, b;
		scanf("%d%d",&a,&b);
		add_edge(a,b);
		add_edge(b,a);
	}
}
 
const int Tree_const = 50; 
 
int root[maxn];
 
struct Chair_Tree{
	int cnt_node;
 
	int sumv[maxn * Tree_const], lson[maxn * Tree_const], rson[maxn * Tree_const];
 
	void build(int l, int r, int &o){
		if(l > r) return ;
		o = ++ cnt_node;
		if(l == r) return ;
		int mid = (l + r) >> 1;
		build(l, mid, lson[o]);
		build(mid + 1, r, rson[o]);
	}
 
	int insert(int l, int r, int o, int pos){
		int oo = ++cnt_node;
		lson[oo] = lson[o];
		rson[oo] = rson[o];
		sumv[oo] = sumv[o] + 1;
 
		if(l == r) return oo;
 
		int mid = (l + r) >> 1;
		if(pos <= mid) lson[oo] = insert(l, mid, lson[o], pos);
		else rson[oo] = insert(mid + 1, r, rson[o], pos);
		return oo;
	}
 
	int query(int l, int r, int u, int v, int lca, int lca_fa, int k){
		if(l == r) return l;
		int lsum = sumv[lson[u]] + sumv[lson[v]] - sumv[lson[lca]] - sumv[lson[lca_fa]];
		int mid = (l + r) >> 1;
		if(k <= lsum) return query(l, mid, lson[u], lson[v], lson[lca], lson[lca_fa], k);
		else return query(mid + 1, r, rson[u], rson[v], rson[lca], rson[lca_fa], k - lsum);
	}
 
}Tree;
 
const int logn = 20;
 
int f[23][maxn];
 
int dep[maxn];
 
void dfs(int u, int fa, int depth){
 
	root[u] = Tree.insert(1, n, root[fa], val[u]);
	dep[u] = depth;
	f[0][u] = fa;
 
	for(int v = head[u]; v ; v = nex[v]){
		if(to[v] == fa) continue;
		dfs(to[v], u, depth + 1);
	}
}
 
inline void get_ancester(){
	for(int i = 1;i <= logn; ++i){
		for(int j = 1;j <= n; ++j)
			f[i][j] = f[i - 1][f[i - 1][j]];
	}
}
 
inline int get_lca(int a, int b){
	if(dep[a] > dep[b]) swap(a,b);
	if(dep[a] != dep[b]){
	    for(int i = logn;i >= 0;--i){
		    if(dep[f[i][b]] >= dep[a]) b = f[i][b];
	    }
	}
	if(a == b) return a;
	for(int i = logn;i>=0;--i)
		if(f[i][a] != f[i][b]) a = f[i][a], b = f[i][b];
	return f[0][a];
}
 
inline void Build(){
	Tree.build(1, n, root[0]);
	dfs(1, 0, 1);
	get_ancester();
}
 
inline void Init(){
	Read();
	Disperse();
	Build();
}
 
inline void Work(){
 
	int lastans = 0;
 
	while(m--){
		int u, v, k;
		scanf("%d%d%d",&u, &v, &k);
	//	u ^= lastans;
 
		int lca = get_lca(u, v);
 
		lastans = Tree.query(1, n, root[u], root[v], root[lca], root[f[0][lca]], k);
		lastans = Sorted[lastans];
 
		printf("%d\n", lastans);
	}
}
 
int main(){
	SetIO("input");
	Init();
	Work();
	return 0;
}