【SPOJ-COT】Count on a tree【主席树】【LCA】

【题目链接】

同【BZOJ2588题解】

除过不强制在线，都是一样的。

/* Pigonometry */
#include 
#include 

using namespace std;

typedef long long LL;

const int maxn = 100005, maxnode = 2000005;

int n, m, tot, num[maxn], rank[maxn], head[maxn], cnt;

struct _edge {
	int v, next;
} g[maxn << 1];

inline int iread() {
	int f = 1; LL x = 0; char ch = getchar();
	for(; ch < '0' || ch > '9'; ch = getchar()) f = ch == '-' ? -1 : 1;
	for(; ch >= '0' && ch <= '9'; ch = getchar()) x = x * 10 + ch - '0';
	return f * x;
}

inline void add(int u, int v) {
	g[cnt] = (_edge){v, head[u]};
	head[u] = cnt++;
}

/* binary search */
inline int binfind(int c) {
	int l = 1, r = tot;
	while(l <= r) {
		int mid = l + r >> 1;
		if(num[mid] <= c) l = mid + 1;
		else r = mid - 1;
	}
	return r;
}

/* fotile tr & doubling lca */
int son[maxnode][2], segcnt, sum[maxnode], root[maxn];
int pre[maxn][18], depth[maxn];

inline void insert(int x, int &y, int l, int r, int c) {
	y = ++segcnt;
	sum[y] = sum[x] + 1;
	if(l == r) return;
	son[y][0] = son[x][0]; son[y][1] = son[x][1];
	int mid = l + r >> 1;
	if(c <= mid) insert(son[x][0], son[y][0], l, mid, c);
	else insert(son[x][1], son[y][1], mid + 1, r, c);
}

inline void dfs(int x) {
	insert(root[pre[x][0]], root[x], 1, tot, rank[x]);
	for(int i = head[x]; ~i; i = g[i].next) if(g[i].v ^ pre[x][0]) {
		pre[g[i].v][0] = x;
		depth[g[i].v] = depth[x] + 1;
		dfs(g[i].v);
	}
}

inline int getlca(int u, int v) {
	if(depth[u] < depth[v]) swap(u, v);
	for(int i = 17; i >= 0; i--) if(depth[pre[u][i]] >= depth[v]) u = pre[u][i];
	for(int i = 17; i >= 0; i--) if(pre[u][i] != pre[v][i]) u = pre[u][i], v = pre[v][i];
	return u == v ? u : pre[u][0];
}

inline int query(int x, int y, int k) {
	int a = x, b = y, c = getlca(x, y), d = pre[c][0];
	a = root[x]; b = root[y]; c = root[c]; d = root[d];
	int l = 1, r = tot;
	while(l < r) {
		int mid = l + r >> 1;
		int tmp = sum[son[a][0]] + sum[son[b][0]] - sum[son[c][0]] - sum[son[d][0]];
		if(tmp >= k) r = mid, a = son[a][0], b = son[b][0], c = son[c][0], d = son[d][0];
		else k -= tmp, l = mid + 1, a = son[a][1], b = son[b][1], c = son[c][1], d = son[d][1];
	}
	return num[l];
}

/* main */
int main() {
	n = iread(); m = iread();
	for(int i = 1; i <= n; i++) num[i] = rank[i] = iread();

	sort(num + 1, num + 1 + n);
	tot = unique(num + 1, num + 1 + n) - num - 1;
	for(int i = 1; i <= n; i++) rank[i] = binfind(rank[i]);

	for(int i = 1; i <= n; i++) head[i] = -1; cnt = 0;
	for(int i = 1; i < n; i++) {
		int u = iread(), v = iread();
		add(u, v); add(v, u);
	}

	depth[1] = 1;
	dfs(1);
	for(int i = 1; i < 18; i++) for(int j = 1; j <= n; j++)
		pre[j][i] = pre[pre[j][i - 1]][i - 1];

	while(m--) {
		int x = iread(), y = iread(), k = iread();
		printf("%d\n", query(x, y, k));
	}
	return 0;
}