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