P3047 附近的牛

题面:https://www.luogu.org/problemnew/show/P3047

Code:
# include <bits/stdc++.h>
using namespace std;
const int root=1,MAXN=100001;
struct rec{
    int pre,to;
}a[2*MAXN];
int head[MAXN],f[MAXN][21],w,n,tot=0,c[MAXN];
bool vis[MAXN];
void adde(int u,int v)
{
    a[++tot].pre=head[u];
    a[tot].to=v;
    head[u]=tot;
}
void dfs1(int u)
{
    vis[u]=true;
    for (int i=head[u];i!=0;i=a[i].pre)
    {
        int v=a[i].to; if (vis[v]) continue;
        dfs1(v);
        for (int j=1;j<=w;j++) f[u][j]+=f[v][j-1];
    }
}
void dfs2(int u)
{
    vis[u]=true;
    for (int i=head[u];i!=0;i=a[i].pre)
    {
        int v=a[i].to;if (vis[v]) continue;
        for (int j=w;j>=2;j--) f[v][j]-=f[v][j-2];
        for (int j=1;j<=w;j++) f[v][j]+=f[u][j-1];
        dfs2(v);
    }
}
int main()
{
    scanf("%d%d",&n,&w);
    int u,v;
    for (int i=1;i<=n-1;i++) {
        scanf("%d%d",&u,&v);
        adde(u,v); adde(v,u);
    }
    memset(f,0,sizeof(f));
    for (int i=1;i<=n;i++) {
        scanf("%d",&c[i]);
        f[i][0]=c[i];
    }
    memset(vis,false,sizeof(vis));
    dfs1(root);
    memset(vis,false,sizeof(vis));
    dfs2(root);
    for (int i=1;i<=n;i++) {
        int ans=0;
        for(int j=0;j<=w;j++) ans+=f[i][j];
        printf("%d\n",ans);
    }
    return 0;
}