[USACO12FEB]Nearby Cows G

题目链接：https://www.luogu.com.cn/problem/P3047

思路：首先要求的是树中所有的子节点 那么考虑换根dp  然后考虑如何通过父节点来更新子节点

虽然k很小 但是暴力更新可能会有很多层所以TLE 考虑dp[i][j] 表示距离以i为根节点距离为j的值

所以每次 换根的时候 只需要把 根节点的减去子节点那一边的dp 就可以得到另外一边的贡献

然后回溯的时候记得 把dp复原即可

#include<bits/stdc++.h>
using namespace std;
const int maxn=1e5+10;
const int mod=998244353;
#define ll long long
#define ull unsigned long long
#define pi pair<int,ll>
#define fi first
#define sc second
#define pb push_back
vector<int>E[maxn];
int c[maxn];
int n,k;
int dp[maxn][30];
int sum[maxn];

void dfs(int u,int fa)
{
    dp[u][0]=c[u];
    for(auto &v:E[u])
    {
        if(v==fa) continue;
        dfs(v,u);
        for(int i=1;i<=k;i++)
            dp[u][i]+=dp[v][i-1];
    }
}
void add(int u,int v,int t)
{
    for(int i=1;i<=k;i++)
    {
        if(t) dp[u][i]+=dp[v][i-1];
        else dp[u][i]-=dp[v][i-1];
    }
}

void dfs22(int u,int fa)
{
    for(int i=0;i<=k;i++) sum[u]+=dp[u][i];
    for(auto &v:E[u])
    {
        if(v==fa) continue;
        add(u,v,0),add(v,u,1);
        dfs22(v,u);
        add(v,u,0),add(u,v,1);
    }
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    cin>>n>>k;
    for(int i=1;i<n;i++)
    {
        int x,y;
        cin>>x>>y;
        E[x].pb(y);
        E[y].pb(x);
    }
    for(int i=1;i<=n;i++) cin>>c[i];
    dfs(1,0);
    dfs22(1,0);
    for(int i=1;i<=n;i++) cout<<sum[i]<<'\n';

}