python/TangDou/Topic/HuanGenDp/P3047.cpp

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

int n, q;
const int N = 200010, K = 21;
int head[N], dp[N][K];
struct ljj {
    int to, stb;
} a[N];
int s;

void insert(int x, int y) {
    s++;
    a[s].stb = head[x];
    a[s].to = y;
    head[x] = s;
}

void dfs1(int x, int fa) {
    for (int i = head[x]; i; i = a[i].stb) {
        int xx = a[i].to;
        if (xx == fa)
            continue;
        dfs1(xx, x);
        for (int j = 1; j <= q; j++)
            dp[x][j] += dp[xx][j - 1]; // 第一遍dp
    }
}

void dfs2(int x, int fa) {
    for (int i = head[x]; i; i = a[i].stb) {
        int xx = a[i].to;
        if (xx == 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[xx][j] -= dp[xx][j - 2]; // 简单容斥
        for (int j = 1; j <= q; j++)
            dp[xx][j] += dp[x][j - 1]; // 第二遍dp
        dfs2(xx, x);
    }
}
signed main() {
    cin >> n >> q;
    for (int i = 1; i < n; i++) {
        int x, y;
        cin >> x >> y;
        insert(x, y);
        insert(y, x);
    }
    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统计答案
        printf("%lld\n", ans);
    }
}