url: https://www.luogu.com.cn/problem/P3047
tag:
树形DP, USACO, 2012
思路：
用f[u,m] 表示以u为根节点，向下走出不超过m步时的权值，d[u, m] 表示以u为根节点，向上或向下走不超过m步时的权值。利用两次dfs分别求出f和d，最后输出每个节点作为根节点时的权值和即可。
代码：
#include <iostream>
#include <algorithm>
#include <cstring>
#include <cstdio>
using namespace std;
typedef long long LL;
const int N = 1e5 + 10;
int h[N], e[N * 2], ne[N * 2], idx;
int w[N];
int n, k;
LL f[N][25], d[N][25];
void add(int a, int b)
{
e[idx] = b;
ne[idx] = h[a];
h[a] = idx ++;
}
void dfs(int u, int fa)
{
for (int i = 0; i <= k; i ++) f[u][i] = w[u];
for (int i = h[u]; ~i; i = ne[i])
{
int j = e[i];
if (j != fa)
{
dfs(j, u);
for (int p = 1; p <= k; p ++)
{
f[u][p] += f[j][p - 1];
}
}
}
}
void dfs2(int u, int fa)
{
for (int i = h[u]; ~i; i = ne[i])
{
int j = e[i];
if (j == fa) continue;
d[j][1] += d[u][0];
for (int p = 2; p <= k; p ++)
{
d[j][p] += d[u][p - 1] - f[j][p - 2];
}
dfs2(j, u);
}
}
int main()
{
cin >> n >> k;
memset(h, -1, sizeof h);
for (int i = 1; i <= n - 1; i ++)
{
int u, v;
cin >> u >> v;
add(u, v), add(v, u);
}
for (int i = 1; i <= n; i ++) cin >> w[i];
dfs(1, 0);
for (int i = 1; i <= n; i ++)
for (int j = 0; j <= k; j ++)
d[i][j] = f[i][j];
dfs2(1, 0);
for (int i = 1; i <= n; i ++) cout << d[i][k] << endl;
return 0;
}