P3047 [USACO12FEB] Nearby Cows G

/*
求距离每一个点 0~k的 值 之和
f[u][i] 距离 u dis==i 的 val和(f[u][]的前缀和很难写)
pre-dfs-> f[u][i]=sum(f[v][i-1])->u的子树内
f[u][i]=f[fa[u]][i-1]+f[u][i]-sum(f[u][0~i-2])
*/
/*
6 2 
5 1 
3 6 
2 4 
2 1 
3 2 
1 
2 
3 
4 
5 
6 

15
21
16
10
8
11

*/
#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cmath>
#include<string.h>
#include<queue>
#include<vector>
#include<bits/stdc++.h>
#define ll long long
#define ddd printf("-----------------------\n");
using namespace std;
const int maxn=1e5 +10;
const int mod=998244353;
const int inf=0x3f3f3f3f;

int head[maxn],to[maxn<<1],nxt[maxn<<1],tot;
int n,k,f[maxn][25],ans,dep[maxn];

void add(int a,int b){
    to[++tot]=b,nxt[tot]=head[a],head[a]=tot;
}

void dfs1(int u,int fa)
{
    dep[u]=dep[fa]+1;
    for(int i=head[u];i;i=nxt[i])
    {
        int v=to[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)
{
    for(int i=head[u];i;i=nxt[i])
    {
        int v=to[i];if(dep[v]<=dep[u]) continue;
        for(int j=k;j>=2;j--)
            f[v][j]-=f[v][j-2];
        for(int j=1;j<=k;j++)
            f[v][j]+=f[u][j-1];
        dfs2(v);
    }
}

int main()
{
    ios::sync_with_stdio(false);
    cin>>n>>k;
    for(int i=1;i<=n-1;i++)
    {
        int a,b;cin>>a>>b;
        add(a,b),add(b,a);
    }
    for(int i=1;i<=n;i++) cin>>f[i][0];
    
    dfs1(1,1);
    dfs2(1);
    
    for(int i=1;i<=n;i++){
        ans=0;
        for(int j=0;j<=k;j++)
            ans+=f[i][j];
        cout<<ans<<'\n';
    }
    
    return 0;
}