P2633 Count on a tree（主席树 + lca）

P2633 Count on a tree

题意

$n$ 个点的树，每个点有一个权值。$m$ 次查询，每次给出 $u,v,k$， 询问 $u,v$ 路径上的点权第 $k$ 小值。强制在线。

思路

本题不需要树链剖分，用树上前缀和就可以。对每个点开一棵权值线段树，维护该点到根节点路径上的权值分布情况。由于每个点都是由根节点的权值线段树上修改一个位置的权值得到，所以自然的可以用主席树维护。

那么怎么查询呢？回顾一下我们怎样通过树上前缀和求两点间点权和？假设每个点到根节点的点权和是 $pre[i]$，那么有：
$$dist(u,v)=pre[u]+pre[v]-pre[lca(u,v)]-pre[fa[lca(u,v)]]$$

对于本题也是同样的，只不过变成了四个权值线段树的加减操作。具体见代码

#include <bits/stdc++.h>
using namespace std;
// #define int long long
#define PII pair<int, int>
#define endl "\n"
/**********************  Core code begins  **********************/
const int N = 5e6 + 7;
struct perseg {
	struct Info {
		int l, r, sum;
	};
	int cnt = 0;
	array<Info, N> info = {};

	int build(int l, int r) {
		int rt = ++cnt;
		int mid = (l + r) >> 1;
		if (l < r) {
			info[rt] = {build(l, mid), build(mid + 1, r), 0};
		}
		return rt;
	}

	int modify(int p, int l, int r, int x, int k) {
		int rt = ++cnt;
		info[rt] = {info[p].l, info[p].r, info[p].sum + k};
		if (l >= r) {
			return rt;
		}
		int mid = (l + r) >> 1;
		if (x <= mid) {
			info[rt].l = modify(info[p].l, l, mid, x, k);
		} else {
			info[rt].r = modify(info[p].r, mid + 1, r, x, k);
		}
		return rt;
	}

	// 分别是 u, v, lca(u, v), fa[lca(u, v)]
	int query(int p1, int p2, int p3, int p4, int l, int r, int k) {
		if (l >= r) {
			return l;
		}
		int x = info[info[p1].l].sum + info[info[p2].l].sum
				- info[info[p3].l].sum - info[info[p4].l].sum;
		int mid = (l + r) >> 1;
		if (x >= k) {
			return query(info[p1].l, info[p2].l, info[p3].l, info[p4].l, l, mid, k);
		} else {
			return query(info[p1].r, info[p2].r, info[p3].r, info[p4].r, mid + 1, r, k - x);
		}
	}
};

void SolveTest() {
	int n, m;
	cin >> n >> m;
	vector<int> a(n + 1);
	for (int i = 1; i <= n; i++) {
		cin >> a[i];
	}

	auto b = a;
	sort(b.begin() + 1, b.end());
	int len = unique(b.begin() + 1, b.end()) - b.begin() - 1;
	auto rank = [&](int x) {
		return lower_bound(b.begin() + 1, b.begin() + 1 + len, x) - b.begin();
	};
	auto kth = [&](int k) {
		return b[k];
	};
	for (int i = 1; i <= n; i++) {
		a[i] = rank(a[i]);
	}

	vector<vector<int>> g(n + 1);
	for (int i = 1; i < n; i++) {
		int u, v;
		cin >> u >> v;
		g[u].push_back(v);
		g[v].push_back(u);
	}

	perseg tr;
	vector<int> root(n + 1), dep(n + 1);
	vector<array<int, 30>> fa(n + 1);

	const int M = 25;

	function<void(int, int)> dfs1 = [&](int u, int pa) {
		root[u] = tr.modify(root[pa], 1, n, a[u], 1);
		dep[u] = dep[pa] + 1;
		fa[u][0] = pa;
		for (int i = 1; i <= M; i++) {
			fa[u][i] = fa[fa[u][i - 1]][i - 1];
		}
		for (int v : g[u]) {
			if (v == pa) {
				continue;
			}
			dfs1(v, u);
		}
	};

	function<int(int, int)> lca = [&](int u, int v) {
        if (dep[u] < dep[v]) {
            swap(u, v);
        }
        int t = dep[u] - dep[v];
        for (int i = 0; i < M; i++) {
            if (t & (1 << i)) {
                u = fa[u][i];
            }
        }
        for (int i = M - 1; i >= 0; i--) {
            if (fa[u][i] != fa[v][i]) {
                u = fa[u][i];
                v = fa[v][i];
            }
        }
        return u == v ? u : fa[u][0];
	};

	root[0] = tr.build(1, n);
	dfs1(1, 0);

	int last = 0;
	for (int i = 1; i <= m; i++) {
		int u, v, k;
		cin >> u >> v >> k;
		u ^= last;
		int l = lca(u, v);
		last = tr.query(root[u], root[v], root[l], root[fa[l][0]], 1, n, k);
		last = kth(last);
		cout << last << endl;
	}
}

/**********************  Core code ends  ***********************/
signed main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	int T = 1;
	// cin >> T;
	for (int i = 1; i <= T; i++) {
		SolveTest();
	}
	return 0;
}