SPOJ - COT Count on a tree（树上主席树 求树中路径第 k 小权值 + 倍增 LCA）

前置知识：普通主席树

题意：

给定 一个包含 N 个结点的树 树节点 从 1 到 N 编号。每个节点有一个 整数权值。

要求执行 以下操作：

u v k：询问 从节点 u 到 节点 v 的路径上 的 第 k 小的权值。

思路：

在普通主席树 求静态数组第 k 小值 时，我们在 for 循环 建树，并利用前缀和的思想每次查询 某段区间的第 k 小值。

普通主席树 维护的其实是一个 线性表的前缀和，但前缀和不一定要出现在线性表上，即：对于 树中 每个 从根到点 u 的路径 是一个 线性序列，我们也可以 把这个序列建成主席树。

显然本题是 求树上任意路径中的第 k 小权值，为了实现上面这点，我们需要引入一个 新概念，树上主席树。

普通区间主席树 和 树上主席树 的区别：

这两个其实只有 建树方式不同 而已，普通主席树在 for 循环里面建树，for 循环 就相当于 线性的区间建树，而 树上主席树是在树形结构上建主席树，建出来的主席树有 树的性质，所以我们要求 任意两点之间最短路径上的第 k 小，就需要用到 lca（至于为什么，后面有解释），在 树上建树 其实就是 在 dfs 遍历树的时候 我们将 当前版本 和 上一版本 进行 复制，修改，而 上一版本 就是 父节点。

树上建主席树 代码：

void dfs(int u, int f)
{
	root[u] = insert(root[f], 0, nums.size() - 1, find(w[u]));	//最新版本：root[u] 上一版本：root[f]
	for (int i = 0; i < g[u].size(); ++i)
	{
		int v = g[u][i];
		if (v == f) continue;
		dfs(v, u);
	}
}

维护 树上任意路径第 k 小权值 的具体操作，结合图例来分析：

（上图中，0 是自己添加的作为 节点 1 的父节点）

现在，我们要 查询 2 到 3 这条路径上的第 k 小权值，

由于在 dfs 建树 的时候由于已经分别保存了 从 0 到 2 和从 0 到 3 的前缀信息，所以 2，3 之间的路径信息 就可以用 t[2].cnt + t[3].cnt - t[lca(2, 3)] - t[fa[lca(2, 3)]] 这个式子来求得，其中 lca(2, 3) 表示 节点 2 和节点 3 的最近公共祖先，fa[lca] 表示 lca 的父节点。

（解释：在 t[2].cnt + t[3].cnt 中，t[lca(2, 3)] 被加了 两遍，因此要 - t[lca(2, 3)]，之后类比 前缀和 的思想 [l, r] = sum[r] - sum[l - 1]，最后要 - t[fa[lca(2, 3)]]）

所以，对于 一般的 u , v 路径之间的信息 就可以用 t[u].cnt + t[v].cnt - t[lca(u, v)] - t[fa[lca(u, v)]].cnt 进行查询。

深搜、递推 预处理 + 树上倍增求 lca 板子（简化版 更好用 更好记）：

const int N = 1e5 + 10;
int n, m;	//n 个点 m 次询问

//LCA
vector<int> g[N];	//存树
int fa[N][18], depth[N];
int LOG;

void dfs(int u, int f)
{
	fa[u][0] = f;
	depth[u] = depth[f] + 1;
	for (int j = 1; j <= LOG; ++j) {
		fa[u][j] = fa[fa[u][j - 1]][j - 1];
	}
	for (int i = 0; i < g[u].size(); ++i)
	{
		int v = g[u][i];
		if (v == f) continue;
		dfs(v, u);
	}
}

int lca(int x, int y)
{
	if (depth[x] < depth[y]) swap(x, y);
	int dc = depth[x] - depth[y];
	for (int k = 0; k <= LOG; ++k) {
		if ((1 << k) & dc) x = fa[x][k];
	}
	if (x == y) return x;
	for (int k = LOG; k >= 0; --k) {
		if (fa[x][k] != fa[y][k]) {
			x = fa[x][k];
			y = fa[y][k];
		}
	}
	x = fa[x][0];
	return x;
}

/*
int main(){
	cin >> n >> m;
	
	//此处完成建树操作：blablabla...
	
	LOG = int(log(n + 0.0) / log(2.0)) + 1;
	dfs(1, 0);
	while (m--)
	{
		int a, b; scanf("%d%d", &a, &b);
		printf("%d\n", lca(a, b));
	}
	return 0;
}
*/

注意：点权 需要开 long long。

时间复杂度：

$O(mlogn)$

代码：

#include <bits/stdc++.h>

using namespace std;
//#define map unordered_map
//#define int long long
#define ll long long
const int N = 1e5 + 10, M = 1e5 + 10;
const int ALL = (N << 2) + M * 17;
int n, m;
vector<int> g[N];

//HJT Tree
ll w[N];
vector<ll> nums;
int root[ALL];

struct node
{
	int l, r;
	int cnt;
} t[ALL];
int idx;

int find(int x) {
	return lower_bound(nums.begin(), nums.end(), x) - nums.begin();
}

int build(int l, int r) {
	int p = ++idx;
	if (l == r) {
		return p;
	}
	int mid = l + r >> 1;
	t[p].l = build(l, mid), t[p].r = build(mid + 1, r);
	return p;
}

int insert(int p, int l, int r, int x)
{
	int q = ++idx;
	t[q] = t[p];
	if (l == r)
	{
		++t[q].cnt;
		return q;
	}
	int mid = l + r >> 1;
	if (x <= mid) t[q].l = insert(t[p].l, l, mid, x);
	else t[q].r = insert(t[q].r, mid + 1, r, x);
	t[q].cnt = t[t[q].l].cnt + t[t[q].r].cnt;
	return q;
}

int ask(int q, int p, int lca, int fa_lca, int l, int r, int x)
{
	if (l == r) return r;
	int cnt = t[t[q].l].cnt + t[t[p].l].cnt - t[t[lca].l].cnt - t[t[fa_lca].l].cnt;
	int mid = l + r >> 1;
	if (x <= cnt) return ask(t[q].l, t[p].l, t[lca].l, t[fa_lca].l, l, mid, x);
	else return ask(t[q].r, t[p].r, t[lca].r, t[fa_lca].r, mid + 1, r, x - cnt);
}

//LCA
int fa[N][18], depth[N];
int LOG;

void dfs(int u, int f)
{
	fa[u][0] = f;
	depth[u] = depth[f] + 1;
	for (int j = 1; j <= LOG; ++j) {
		fa[u][j] = fa[fa[u][j - 1]][j - 1];
	}
	root[u] = insert(root[f], 0, nums.size() - 1, find(w[u]));
	for (int i = 0; i < g[u].size(); ++i)
	{
		int v = g[u][i];
		if (v == f) continue;
		dfs(v, u);
	}
}

int lca(int x, int y)
{
	if (depth[x] < depth[y]) swap(x, y);
	int dc = depth[x] - depth[y];
	for (int k = 0; k <= LOG; ++k) {
		if ((1 << k) & dc) x = fa[x][k];
	}
	if (x == y) return x;
	for (int k = LOG; k >= 0; --k) {
		if (fa[x][k] != fa[y][k]) {
			x = fa[x][k];
			y = fa[y][k];
		}
	}
	x = fa[x][0];
	return x;
}

signed main()
{
	cin >> n >> m;
	for (int i = 1; i <= n; ++i)
	{
		scanf("%lld", &w[i]);
		nums.push_back(w[i]);
	}
	int t = n - 1;
	while (t--)
	{
		int u, v; scanf("%d%d", &u, &v);
		g[u].push_back(v);
		g[v].push_back(u);
	}

	sort(nums.begin(), nums.end());
	nums.erase(unique(nums.begin(), nums.end()), nums.end());

	root[0] = build(0, nums.size() - 1);

	LOG = int(log(n + 0.0) / log(2.0)) + 1;
	dfs(1, 0);	//lca深搜预处理 与此同时建树上主席树

	while (m--)
	{
		int u, v, k;
		scanf("%d%d%d", &u, &v, &k);
		int LCA = lca(u, v);
		int ans = ask(root[v], root[u], root[LCA], root[fa[LCA][0]], 0, nums.size() - 1, k);
		printf("%lld\n", nums[ans]);
	}

	return 0;
}