python/TangDou/Topic/HuanGenDp/P3047.cpp

#include <bits/stdc++.h>
using namespace std;
const int N = 1e5 + 10, M = N << 1;
const int K = 25;
// 链式前向星
int e[M], h[N], idx, w[M], ne[M];
void add(int a, int b, int c = 0) {
    e[idx] = b, ne[idx] = h[a], w[idx] = c, h[a] = idx++;
}

int f[N][K]; // f[i][j]:如果根是1号节点时，i号节点，最远走j步，可以获取到的所有点权和
int g[N][K];
int val[N]; // 点权数组

int n, k;

void dfs1(int u, int fa) {
    for (int i = 0; i <= k; i++) f[u][i] = val[u];

    for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v == fa) continue;
        dfs1(v, u);

        for (int j = 1; j <= k; j++) f[u][j] += f[v][j - 1];
    }
}

void dfs2(int u, int fa) {
    for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v == fa) continue;

        g[v][0] = val[v];
        g[v][1] = f[v][1] + val[u];

        for (int j = 2; j <= k; j++) g[v][j] = g[u][j - 1] + f[v][j] - f[v][j - 2];
        dfs2(v, u);
    }
}

int main() {
    // 初始化链式前向星
    memset(h, -1, sizeof h);

    cin >> n >> k;
    for (int i = 1; i < n; i++) { // n-1条边
        int a, b;
        cin >> a >> b;
        add(a, b), add(b, a);
    }

    for (int i = 1; i <= n; i++) cin >> val[i]; // 点权
    // 1、自底向上
    dfs1(1, 0);

    // 2、换根dp
    for (int i = 0; i <= k; i++) g[1][i] = f[1][i];
    dfs2(1, 0);

    // 输出结果
    for (int i = 1; i <= n; i++) cout << g[i][k] << endl;
    return 0;
}