计蒜客 ICPC NEAU Programming Contest 2020 G.选根 （树形dp）

链接 ：选根

题意：
定义根的深度为 1， 节点 i 的价值为 depth[i]*val[i]， 求以 1 -n 为根时 整棵树的价值。

思路
这种题目一般都是分成子树内和子树外分别求值的。

1. 首先来求子树内的权值， dp[u] 表示 以 u 为根的子树内的权值 ，vv[u]表示以 u 的子树内所有点的 val 之和。那么转移方程就是 dp[u]+=dp[to]+vv[to]; vv[u]+=vv[to];
   因为向父节点转移的时候相当于 每个节点的深度都 +1 了 所以要加上 vv[to];
2. 然后就是 子树外的权值，子树外的权值可以用父节点的答案减去 当前节点对父节点的贡献，再加上子树外的节点深度+1 产生的贡献，再加上该节点为根 子树内的贡献就好了。

pv[u]=vv[pre]+pv[pre]-vv[u];   // pv 表示子树外 点的 val值之和
ans[u]=ans[pre]-dp[u]-vv[u] +pv[u] +dp[u];  dp[u]可抵消，这样写便于理解

代码

#include<iostream>
#include<cstdio>
#include<set>
#include<math.h>
#include<queue>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
const int maxn=4e5+7;
int n,m,head[maxn],num=0;
ll dp[maxn],pre[maxn],ans[maxn],val[maxn],vv[maxn],pv[maxn],av[maxn];
struct node{
    ll next,to;
}e[maxn];
void add(int u,int v){
     e[num].next=head[u];
     e[num].to=v;
     head[u]=num++;
}
void dfs1(int u, int pre){
    dp[u]=val[u];vv[u]=val[u];
    for(int i=head[u];i!=-1;i=e[i].next){
        int to=e[i].to;
        if(to== pre)  continue;
        dfs1(to,u);
        dp[u]+=dp[to]+vv[to];
        vv[u]+=vv[to];
    }
}
void dfs2(int u, int pre){
    if(u==1) pv[u]=0,ans[u]=dp[u];
    else {
        pv[u]=vv[pre]+pv[pre]-vv[u];
        ans[u]=ans[pre]-dp[u]-vv[u] +pv[u] +dp[u];
    }
    for(int i=head[u];i!=-1;i=e[i].next){
        int to=e[i].to;
        if(to== pre)  continue;
        dfs2(to,u);
    }
}
int main (){
   int t;
   cin>>t;
   while(t--){
        int n;cin>>n;
        for(int i=0;i<=n;i++) head[i]=-1;num=0;
        for(int i=1;i<=n;i++) scanf("%lld",&val[i]);
        for(int i=0,u,v;i<n-1;i++){
            scanf("%d%d",&u,&v);
            add(u,v);
            add(v,u);
        }
        dfs1(1,-1);dfs2(1,-1);
        for(int i=1;i<=n;i++){
            printf ("%lld ",ans[i]);
        }
        printf ("\n");
   }
}