P3047 [USACO12FEB]Nearby Cows G
题目描述
思路
使用换根DP，
设 \(dp[i][j]\) 表示以 \(i\) 为根节点的子树中深度小于等于 \(j\) 的点的权值之和。
设 \(f[i][j]\) 表示将第 \(i\) 个点作为整棵树的根节点深度小于等于 \(j\) 的点的权值之和。
有：
\[\begin{cases}
dp[u][k] = \sum dp[to][k - 1] \\
f[to][j] = dp[to][j] - dp[to][j - 2] + f[u][j - 1]
\end{cases}
\]
\(f[to][j] = dp[to][j] - dp[to][j - 2] + f[u][j - 1]\) 表示：
代码：
#include <bits/stdc++.h>
using namespace std;
const int N = 100010, M = 30;
int w[N];
struct Edge {
int to;
int next;
}e[N * 2];
int head[N], idx;
void add(int a, int b) {
idx++;
e[idx].to = b;
e[idx].next = head[a];
head[a] = idx;
}
int n, k;
int dp[N][M];
int f[N][M];
void dfs(int u, int fa) {
for (int i = 0; i <= k; i++) dp[u][i] = w[u];
for (int i = head[u]; i; i = e[i].next) {
int to = e[i].to;
if (to == fa) continue;
dfs(to, u);
for (int i = 1; i <= k; i++) dp[u][i] += dp[to][i - 1];
}
}
void dfs2(int u, int fa) {
for (int i = head[u]; i; i = e[i].next) {
int to = e[i].to;
if (to == fa) continue;
f[to][1] = dp[to][1] + w[u];
for (int j = k; j >= 2; j--) f[to][j] = f[u][j - 1] - dp[to][j - 2] + dp[to][j];
dfs2(to, u);
}
}
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
cin >> n >> k;
for (int i = 1; i < n; i++) {
int a, b;
cin >> a >> b;
add(a, b);
add(b, a);
}
for (int i = 1; i <= n; i++) cin >> w[i];
dfs(1, 0);
memcpy(f[1], dp[1], sizeof(f[1]));
dfs2(1, 0);
for (int i = 1; i <= n; i++) cout << f[i][k] << '\n';
return 0;
}