[USACO12FEB] 附近的牛 Nearby Cows

给你一棵 (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
",ans);
    }
}