Count on a tree SPOJ - COT LCA+主席树

#include<iostream>
#include<vector>
#include<algorithm>
#include<string>
#include<stack>
#include<cstring>
#include<set>
#include<iterator>
#include<list>
#include<deque>
#include<queue>
#include<map>
#include<cmath>
#include<sstream>
#include<cstdio>
#include<ctime>
#include<iomanip>
#include<unordered_map>
using namespace std;
typedef long long ll;
typedef pair<int, int> P;
const int N = 2e5 + 10;
const int M = N * 30;
const int inf = 0x3f3f3f3f;
const ll INF = 0x3f3f3f3f3f3f3f3f;
int rmq[2 * N];
int n, q, m, tot;
int a[N];
struct segmenttree
{
	int lc, rc;
	int dat;
}t[M];
int root[N];
vector<int> v;
int getid(int x)
{
	return lower_bound(v.begin(), v.end(), x) - v.begin() + 1;
}
int build(int l, int r)
{
	int p = ++tot;
	t[p].dat = 0;
	if (l == r)
		return p;
	int mid = (l + r) >> 1;
	t[p].lc = build(l, mid);
	t[p].rc = build(mid + 1, r);
	return p;
}
int update(int now, int l, int r, int x)
{
	int p = ++tot;
	t[p] = t[now];
	t[p].dat++;
	//cout << t[p].dat << endl;
	if (l == r)
	{
		//cout << "end" << endl;
		return p;
	}
	int mid = (l + r) >> 1;
	if (x <= mid)
	{
		//cout << "left" << endl;
		t[p].lc = update(t[now].lc, l, mid, x);
	}
	else
	{
		//cout << "right" << endl;
		t[p].rc = update(t[now].rc, mid + 1, r, x);
	}
	return p;
}
struct ST
{
	int mm[2 * N];
	int dp[2 * N][20];
	void init(int n)
	{
		mm[0] = -1;
		for (int i = 1; i <= n; i++)
		{
			mm[i] = ((i&(i - 1)) == 0) ? mm[i - 1] + 1 : mm[i - 1];
			dp[i][0] = i;
		}
		for (int j = 1; j <= mm[n]; j++)
		{
			for (int i = 1; i + (1 << j) - 1 <= n; i++)
			{
				dp[i][j] = rmq[dp[i][j - 1]] < rmq[dp[i + (1 << (j - 1))][j - 1]] ?dp[i][j - 1] : dp[i + (1 << (j - 1))][j - 1];
				//printf("dp[%d][%d] == %d\n", i, j, dp[i][j]);
			}
		}
	}
	int query(int a, int b)
	{
		if (a > b)
			swap(a, b);
		int k = mm[b - a + 1];
		return rmq[dp[a][k]] <= rmq[dp[b - (1 << k) + 1][k]] ? dp[a][k] : dp[b - (1 << k) + 1][k];
	}
}st;
int head[N], cnt, ver[2 * N], Next[2 * N];
void addedge(int x, int y)
{
	ver[++cnt] = y, Next[cnt] = head[x], head[x] = cnt;
	ver[++cnt] = x, Next[cnt] = head[y], head[y] = cnt;
}
void init()
{
	memset(head, 0, sizeof(head));
	cnt = 0;
}
int id[N * 2];
int fir[N];
int cnt2;
void dfs(int u, int pre,int dep)
{
	id[++cnt2] = u;
	rmq[cnt2] = dep;
	fir[u] = cnt2;
	for (int i = head[u]; i; i = Next[i])
	{
		int v = ver[i];
		if (v == pre)
			continue;
		dfs(v, u, dep + 1);
		id[++cnt2] = u;
		rmq[cnt2] = dep;
	}
}
void Lca_init(int root, int n)
{
	cnt2 = 0;
	dfs(root, root, 0);
	st.init(2 * n - 1);
}
int query_lca(int u, int v)
{
	return id[st.query(fir[u], fir[v])];
}
void dfs_build(int u, int pre)
{
	//printf("u == %d\n", u);
	int pos = getid(a[u]);
	root[u] = update(root[pre], 1, m, pos);
	//printf("root[%d] == %d ", u, root[u]);
	//cout << t[root[u]].dat << endl;
	for (int i = head[u]; i; i = Next[i])
	{
		int v = ver[i];
		if (v == pre)
			continue;
		dfs_build(v, u);
	}
}
int pos;
int query(int p, int q, int rootlca, int l, int r, int k)
{
	//cout << l << " " << r << endl;
	if (l == r)
		return l;
	int mid = (l + r) >> 1;
//	cout << t[p].dat << " " << t[q].dat << " " << t[rootlca].dat << endl;
	//cout << t[t[p].lc].dat << " " << t[t[q].lc].dat << " " << t[t[rootlca].lc].dat << endl;
	int tmp = t[t[p].lc].dat + t[t[q].lc].dat - 2 * t[t[rootlca].lc].dat + (l <= pos && pos <= mid);
	//cout << "tmp == " << tmp << endl;
	if (tmp >= k)
	{
		return query(t[p].lc, t[q].lc, t[rootlca].lc, l, mid, k);
	}
	else
	{
		return query(t[p].rc, t[q].rc, t[rootlca].rc, mid + 1, r, k - tmp);
	}
}
int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0), cout.tie(0);
	while (cin >> n >> q)
	{
		v.clear();
		init();
		for (int i = 1; i <= n; i++)
		{
			cin >> a[i];
			v.push_back(a[i]);
		}
		sort(v.begin(), v.end());
		v.erase(unique(v.begin(), v.end()), v.end());
		m = v.size();
		tot = 0;		
		int u, vv;
		for (int i = 1; i < n; i++)
		{
			cin >> u >> vv;
			addedge(u, vv);
		}
		root[0] = build(1, m);
		dfs_build(1, 0);
		Lca_init(1, n);
		int k;
		/*for (int i = 1; i <= cnt2; i++)
		{
			printf(" %d ", id[i]);
			printf("rmq[%d] == %d \n", i, rmq[i]);
		}
		cout << endl;*/
		while (q--)
		{
			cin >> u >> vv >> k;
			int lca = query_lca(u, vv);
			pos = getid(a[lca]);
			//cout << pos << endl;
			cout << v[query(root[u], root[vv], root[lca], 1, m, k) - 1] << endl;
		}
	}
	return 0;
}