LG P2633 Count on a tree

( ext{Solution})

树上主席树板子

( ext{Code})

#include <cstdio> 
#include <algorithm>
#define re register
using namespace std;

const int N = 1e5 + 5;
int n, m, f[N][19], len, size, rt[N], a[N], b[N], c[N], tot, h[N], dep[N];
struct segment{int ls, rs, s;}seg[N * 23];
struct edge{int to, nxt;}e[N * 2];
inline void add(int x, int y){e[++tot] = edge{y, h[x]}, h[x] = tot;}

void insert(int &p, int pre, int l, int r, int x)
{
	p = ++size;
	seg[p] = seg[pre], seg[p].s += 1;
	if (l == r) return;
	int mid = l + r >> 1;
	if (x <= mid) insert(seg[p].ls, seg[pre].ls, l, mid, x);
	else insert(seg[p].rs, seg[pre].rs, mid + 1, r, x);
}

void dfs(int x, int fa)
{
	insert(rt[x], rt[fa], 1, len, c[x]);
	for(re int i = 1; i <= 17; i++)
	if (f[x][i - 1]) f[x][i] = f[f[x][i - 1]][i - 1];
	else break;
	for(re int i = h[x]; i; i = e[i].nxt)
	{
		int v = e[i].to;
		if (v == fa) continue;
		f[v][0] = x, dep[v] = dep[x] + 1, dfs(v, x);
	}
}

inline int LCA(int u, int v)
{
	if (dep[u] < dep[v]) swap(u, v);
	int deep = dep[u] - dep[v];
	for(re int i = 0; i <= 17; i++)
	if ((deep >> i) & 1) u = f[u][i];
	if (u == v) return u;
	for(re int i = 17; i >= 0; i--)
	if (f[u][i] ^ f[v][i]) u = f[u][i], v = f[v][i];
	return f[u][0];
}

int query(int u, int v, int x, int y, int l, int r, int k)
{
	if (l == r) return l;
	int s = seg[seg[u].ls].s + seg[seg[v].ls].s - seg[seg[x].ls].s - seg[seg[y].ls].s, mid = l + r >> 1;
	if (k <= s) return query(seg[u].ls, seg[v].ls, seg[x].ls, seg[y].ls, l, mid, k);
	return query(seg[u].rs, seg[v].rs, seg[x].rs, seg[y].rs, mid + 1, r, k - s);
}

inline int querypath(int u, int v, int k)
{
	int x = LCA(u, v), y = f[x][0];
	return b[query(rt[u], rt[v], rt[x], rt[y], 1, len, k)];
}

int main()
{
	scanf("%d%d", &n, &m);
	for(re int i = 1; i <= n; i++) scanf("%d", &a[i]), b[i] = a[i];
	for(re int i = 1, x, y; i < n; i++) scanf("%d%d", &x, &y), add(x, y), add(y, x);
	sort(b + 1, b + n + 1), len = unique(b + 1, b + n + 1) - b;
	for(re int i = 1; i <= n; i++) c[i] = lower_bound(b + 1, b + len + 1, a[i]) - b; dfs(1, 0);
	for(int u, v, k, lst = 0; m; --m) scanf("%d%d%d", &u, &v, &k), printf("%d
", lst = querypath(u ^ lst, v, k));
}