python/TangDou/Topic/HuanGenDp/P3047.cpp

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define endl "\n"

const int N = 200010, M = N << 1, K = 21;

// 链式前向星
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 n, q;
int dp[N][K];

void dfs1(int u, int fa) {
    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 <= q; j++) dp[u][j] += dp[v][j - 1]; // 第一遍dp
    }
}

void dfs2(int u, int fa) {
    for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v == fa) continue;
        // 在第一次遍历时 dp[1][2] 包括了 dp[2][1] 2的子树权值；
        // 然鹅  ans在统计dp[2][3] 的时候也加上了 dp[2][1]  2的子树权值；
        // 第二次遍历 dp[2][3] 又加上了 dp[2][1];
        // 所以需要简单容斥一下；
        for (int j = q; j >= 2; j--)
            dp[v][j] -= dp[v][j - 2]; // 简单容斥
        for (int j = 1; j <= q; j++)
            dp[v][j] += dp[u][j - 1]; // 第二遍dp
        dfs2(v, u);
    }
}
signed main() {
    // 初始化链式前向星
    memset(h, -1, sizeof h);

    cin >> n >> q;
    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 >> dp[i][0]; // 每个节点往外0距离，就是它本身的权值；
    dfs1(1, 0);
    dfs2(1, 0);

    for (int i = 1; i <= n; i++) {
        int ans = 0;
        for (int j = 0; j <= q; j++) ans += dp[i][j]; // ans统计答案
        cout << ans << endl;
    }
}