P3047 [USACO12FEB]Nearby Cows G - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

题意：

思路：

做法就是换根

预处理dp[v][j]用普通的树形DP处理即可

注意：一开始预处理的dp[v][j]指的是在v的子树里离v为j的权值和

Code：

#include <bits/stdc++.h>

#define int long long
#define max(a,b) (a>b?a:b)
#define min(a,b) (a<b?a:b)

using namespace std;

const int mxn=3e6+10;
const int mxe=3e5+10;
const int mod=1e9+7;

struct ty{
    int to,next;
}edge[mxe<<2];

int N,K,u,v,x,tot=0;
int head[mxn],dp[mxn][22];

void add(int u,int v){
    edge[tot].to=v;
    edge[tot].next=head[u];
    head[u]=tot++;
}
void G_init(){
    tot=0;
    for(int i=0;i<=N;i++){
        head[i]=-1;
    }
}
void dfs1(int u,int fa){
    for(int i=head[u];~i;i=edge[i].next){
        if(edge[i].to==fa) continue;
        dfs1(edge[i].to,u);
        for(int j=1;j<=K;j++){
            dp[u][j]+=dp[edge[i].to][j-1];
        }
    }
}
void dfs2(int u,int fa){
    for(int i=head[u];~i;i=edge[i].next){
        if(edge[i].to==fa) continue;
        for(int j=K;j>=2;j--){
            dp[edge[i].to][j]-=dp[edge[i].to][j-2];
        }
        for(int j=1;j<=K;j++){
            dp[edge[i].to][j]+=dp[u][j-1];
        }
        dfs2(edge[i].to,u);
    }
}
void solve(){
    cin>>N>>K;
    G_init();
    for(int i=1;i<=N-1;i++){
        cin>>u>>v;
        add(u,v);
        add(v,u);
    }
    for(int i=1;i<=N;i++){
        cin>>x;
        dp[i][0]=x;
    }
    dfs1(1,0);
    dfs2(1,0);
    for(int i=1;i<=N;i++){
        int ans=0;
        for(int j=0;j<=K;j++) ans+=dp[i][j];
        cout<<ans<<'\n';
    }
}
signed main(){
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    int __=1;//cin>>__;
    while(__--)solve();return 0;
}