题意:有一棵 $n$ 个点的树,问树上距离 $k$ 的点对是否存在。
对于树上一点,它子树中的路径可分为两种:
- 经过它的。
- 以它为端点的。
- 以它为断点分为两条路径。
- 不经过它的。
对于树上一点,只关注它的第一种情况,第二种情况可以继续往下递归求解。
这里设置一个标记数组 $flag$,其中 $flag_i$ 表示以当前节点 $u$ 为根,之前遍历过的子树中是否有到 $u$ 长度为 $i$ 的点。
那么如果当前遍历的子树中,存在 $flag_{dis(u,v)} == 1 (v \in tree_u)$,则存在长度为 $k$ 的路径。
为了保证时间复杂度,每一次递归需要选择子树的重心作为下一次的根节点。
注意区分,重心是用子树大小,计算答案是用边权。
代码:
//by Conan15
#include <bits/stdc++.h>
using namespace std;
const int N = 2e4 + 15, M = N << 1, INF = 20000000;
int n, q, h[N], e[M], w[M], ne[M], idx = 0;
void add(int a, int b, int c) {
e[idx] = b, w[idx] = c, ne[idx] = h[a], h[a] = idx++;
}
bool st[N];
int rt, Max[N], sz[N], sum = 0; //求重心
int query[N];
bool ans[N];
void getsz(int u, int father) {
sz[u] = 1;
Max[u] = 0;
for (int i = h[u]; ~i; i = ne[i]) {
int v = e[i];
if (v == father || st[v]) continue;
getsz(v, u); sz[u] += sz[v];
Max[u] = max(Max[u], sz[v]);
}
Max[u] = max(Max[u], sum - sz[u]);
if (!rt || Max[u] < Max[rt]) rt = u;
}
int d[N];
int p[N], tot = 0;
bool flag[20000015];
void dis(int u, int father) {
p[++tot] = d[u];
for (int i = h[u]; ~i; i = ne[i]) {
int v = e[i];
if (v == father || st[v]) continue;
d[v] = d[u] + w[i];
dis(v, u);
}
}
queue<int> Q;
void dfs(int u, int father) {
st[u] = 1; flag[0] = 1; Q.push(0);
for (int i = h[u]; ~i; i = ne[i]) {
int v = e[i];
if (v == father || st[v]) continue;
d[v] = w[i]; //这步很重要
tot = 0; dis(v, u); //准备求答案
for (int j = 1; j <= q; j++)
for (int k = 1; k <= tot; k++)
if (query[j] >= p[k]) ans[j] |= flag[query[j] - p[k]];
for (int j = 1; j <= tot; j++)
if (p[j] < INF) Q.push(p[j]), flag[p[j]] = 1;
}
while (Q.size()) flag[Q.front()] = 0, Q.pop();
for (int i = h[u]; ~i; i = ne[i]) {
int v = e[i];
if (v == father || st[v]) continue;
rt = 0, sum = sz[v];
getsz(v, u), getsz(rt, -1);
dfs(rt, u);
}
}
int main() {
scanf("%d%d", &n, &q);
for (int i = 1; i <= n; i++) h[i] = -1;
for (int i = 1, a, b, c; i < n; i++) {
scanf("%d%d%d", &a, &b, &c);
add(a, b, c); add(b, a, c);
}
for (int i = 1; i <= q; i++) scanf("%d", &query[i]);
sum = n; rt = 0; getsz(1, -1); getsz(rt, -1);
dfs(rt, -1);
for (int i = 1; i <= q; i++)
puts(ans[i] ? "AYE" : "NAY");
return 0;
}