[USACO12FEB] 附近的牛 Nearby Cows - 树形dp,容斥

给你一棵 \(n\) 个点的树，点带权，对于每个节点求出距离它不超过 \(k\) 的所有节点权值和 \(m_i\)

随便定一个根，设\(f[i][j]\)表示只考虑子树，距离为\(j\)的权值和，\(g[i][j]\)表示考虑子树和父树，距离为\(j\)的权值和，显然答案可以用\(g\)表示

\(f[p][0]=w[p]\)
\(f[p][k]=\sum f[q][k-1]\)
\(g[1][k]=f[1][k]\)
\(g[p][0]=w[p]\)

对\(g\)的计算，考虑容斥

\[g[q][k] =\sum( f[q][k] + g[p][k-1] - f[q][k-2]) \]

注意特判掉\(k=1\)

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 100005;
vector <int> G[N];
int n,k,vis[N],f[N][22],g[N][22],w[N];

void dfs1(int p) {
    vis[p]=1;
    for(int i=0;i<G[p].size();i++) {
        int q=G[p][i];
        if(vis[q]) continue;
        dfs1(q);
        for(int j=1;j<=k;j++)
            f[p][j]+=f[q][j-1];
    }
}

void dfs2(int p) {
    vis[p]=1;
    for(int i=0;i<G[p].size();i++) {
        int q=G[p][i];
        if(vis[q]) continue;
        for(int j=1;j<=k;j++)
            g[q][j]+=f[q][j]+g[p][j-1]-(j>1?1:0)*f[q][j-2];
        dfs2(q);
    }
}

signed main() {
    scanf("%lld%lld",&n,&k);
    for(int i=1;i<n;i++) {
        int t1,t2;
        scanf("%lld%lld",&t1,&t2);
        G[t1].push_back(t2);
        G[t2].push_back(t1);
    }
    for(int i=1;i<=n;i++) scanf("%lld",&w[i]);
    for(int i=1;i<=n;i++) f[i][0]=g[i][0]=w[i];
    dfs1(1);
    for(int i=1;i<=k;i++) g[1][i]=f[1][i];
    memset(vis,0,sizeof vis);
    dfs2(1);
    for(int i=1;i<=n;i++) {
        int ans = 0;
        for(int j=0;j<=k;j++) ans+=g[i][j];
        printf("%lld\n",ans);
    }
}