题意描述
洛谷链接
给你一棵 \(n\)
个点的树，点带权，对于每个节点求出距离它不超过 \(k\) 的所有节点权值和 \(m_i\)。
解题思路
本题树形 DP 可做。首先我们统计 \(f_{i,
j}\)，\(f_{i, j}\) 表示以 \(i\) 为根的子树中距离为 \(j\) 的节点个数。
于是有下列方程：
\[
f_{i, j} = f_{i, j} + f_{e, j - 1} \quad i \rightarrow e
\]
接下来就是换根。对于 \(i\)
节点，我们有一下方程：
\[
f_{i, j} = f_{i, j} + f_{e, j - 1} \quad e \rightarrow i
\]
然而此时必有重复，因为在统计 \(f_{e,
j}\) 时，\(f_{i, j - 2}\)
已被统计。故这里简单容斥一下即可。
\[
f_{i, j} \cup f_{e, j - 1} = f_{i, j} + f_{e, j - 1} - f_{i, j} \cap
f_{e, j - 1}
\]
\[
f_{i, j} = f_{i, j} + f_{e, j - 1} - f_{i, j - 2} \quad e \rightarrow i
\]
代码演示
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
| #include <cstdio>
#include <vector>
const int MAXN = 1e5;
const int MAXK = 20;
struct Node {
std::vector<struct Edge> e;
int f[MAXK + 1];
} N[MAXN + 1];
struct Edge {
Node *s, *t;
Edge(Node *s, Node *t) : s(s), t(t) {}
};
int n, k;
inline void addEdge(int s, int t) {
N[s].e.push_back(Edge(&N[s], &N[t]));
N[t].e.push_back(Edge(&N[t], &N[s]));
}
void dp1(Node *v, Node *fa = nullptr) {
for (Edge *e = &v->e.front(); e && e <= &v->e.back(); e++) {
if (e->t == fa) continue;
dp1(e->t, v);
for (int i = 1; i <= k; i++) v->f[i] += e->t->f[i - 1];
}
}
void dp2(Node *v, Node *fa = nullptr) {
for (Edge *e = &v->e.front(); e && e <= &v->e.back(); e++) {
if (e->t == fa) continue;
for (int i = k; i >= 2; i--) e->t->f[i] -= e->t->f[i - 2];
for (int i = 1; i <= k; i++) e->t->f[i] += v->f[i - 1];
dp2(e->t, v);
}
}
int main() {
scanf("%d %d", &n, &k);
for (int i = 1; i <= n - 1; i++) {
int s, t;
scanf("%d %d", &s, &t);
addEdge(s, t);
}
for (int i = 1; i <= n; i++) scanf("%d", &N[i].f[0]);
dp1(&N[1]);
dp2(&N[1]);
for (int i = 1; i <= n; i++) {
int ans = 0;
for (int j = 0; j <= k; j++) ans += N[i].f[j];
printf("%d\n", ans);
}
return 0;
}
|