「SPOJ COT」 Count on a tree

Posted 2021-02-22 -wallace-

Description

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

执行以下操作 (m) 次：

u v k : 询问从节点 (u) 到 节点 (v) 的路径（包括端点）上的第 (k) 小的权值。

Hint

(1le n, mle 10^5)

Solution

求 Kth ，显然主席树。

问题是怎么放到树上。

设 (T_x) 为从根节点（暂定为 1 ，下同）到结点 (x) 的路径（包括端点）中权值的“前缀主席树”，即 (T_x) 维护这条路径上的权值的集合。

那么我们用一下前缀和的思想，若要得到路径 (u ightarrow v) 的主席树，那么只要：

[T_{u ightarrow v} = T_u + T_v - T_{ ext{LCA}(u, v)} - T_{ ext{father}( ext{LCA}(u, v))} ]

就行了。但实际上不用先把整个 (T_{u ightarrow v}) 求出来，到某个结点的位置时临时计算即可。

LCA 什么的可以上树剖或倍增（下面的代码用了树剖）。

• 时间复杂度 (O((n+m)log n))
• 空间复杂度 (O(nlog n))

Code

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

const int N = 1e5 + 5;
const int S = N << 5;

int lc[S], rc[S], sum[S], total = 0;
int root[N], wei[N];
int n, q;
vector<int> G[N];

#define mid ((l + r) >> 1)
int build(int l, int r) {
	int rt = ++ total;
	if (l == r) return rt;
	lc[rt] = build(l, mid);
	rc[rt] = build(mid + 1, r);
	return rt;
}
int update(int pre, int pos, int l, int r) {
	int rt = ++ total;
	lc[rt] = lc[pre], rc[rt] = rc[pre];
	sum[rt] = sum[pre] + 1;
	if (l == r) return rt;
	if (pos <= mid) lc[rt] = update(lc[pre], pos, l, mid);
	else rc[rt] = update(rc[pre], pos, mid + 1, r);
	return rt;
}
int findKth(int a, int b, int c, int d, int k, int l, int r) {
	if (l == r) return l;
	int x = sum[lc[a]] + sum[lc[b]] - sum[lc[c]] - sum[lc[d]];
	if (k <= x) return findKth(lc[a], lc[b], lc[c], lc[d], k, l, mid);
	else return findKth(rc[a], rc[b], rc[c], rc[d], k - x, mid + 1, r);
}
#undef mid

int fa[N], size[N];
int dep[N], maxs[N];
int top[N];

void Dfs1(int rt, int f) {
	root[rt] = update(root[f], wei[rt], 1, n);
	size[rt] = 1, fa[rt] = f, dep[rt] = dep[f] + 1;
	for (vector<int>::iterator it = G[rt].begin(); it != G[rt].end(); ++ it) {
		if (*it == f) continue;
		Dfs1(*it, rt), size[rt] += size[*it];
		if (size[maxs[rt]] < size[*it])
			maxs[rt] = *it;
	}
}
void Dfs2(int rt,int tp) {
	top[rt] = tp;
	if (maxs[rt]) Dfs2(maxs[rt], tp);
	for (vector<int>::iterator it = G[rt].begin(); it != G[rt].end(); ++ it)
		if (*it != fa[rt] && *it != maxs[rt])
			Dfs2(*it, *it);
}
inline int findLCA(int u,int v) {
	while(top[u] != top[v])
		if (dep[top[u]] > dep[top[v]]) u = fa[top[u]];
		else v = fa[top[v]];
	return dep[u] < dep[v] ? u : v;
}

int buf[N], cnt;
signed main() {
	ios::sync_with_stdio(false);
	cin >> n >> q;
	for (register int i = 1; i <= n; ++ i)
		cin >> wei[i], buf[i] = wei[i];
	for (register int u, v, i = 1; i < n; ++ i) {
		cin >> u >> v;
		G[u].push_back(v);
		G[v].push_back(u);
	}
	
	sort(buf + 1, buf + 1 + n), cnt = unique(buf + 1, buf + 1 + n) - buf - 1;
	for (register int i = 1; i <= n; i ++)
		wei[i] = lower_bound(buf + 1, buf + 1 + cnt, wei[i]) - buf;
		
	root[0] = build(1, n);
	Dfs1(1, 0), Dfs2(1, 1);
	
	while (q --) {
		int u, v, k, l, f;
		cin >> u >> v >> k;
		l = findLCA(u, v), f = fa[l];
		cout << buf[findKth(root[u], root[v], root[f], root[l], k, 1, n)] << endl;
	}
	return 0;
}