Codeforces 1092F Tree with Maximum Cost（树形DP）

题目链接：Tree with Maximum Cost

题意：给定一棵树，树上每个顶点都有属性值ai，树的边权为1，求$sumlimits_{i = 1}^{n} dist(i, v) cdot a_i$，$dist(i, v) $为顶点i到顶点v的距离。该顶点v可以任意选择。

题解：O（n^2）的做法：从每个顶点跑一遍DFS，计算贡献值，并更新答案。（超时）

我们可以先计算出从顶点1跑的答案，发现顶点之间贡献的转移为$ans[u]=ans[fa]+(all-sum[u])-sum[u]$。（all为$sumlimits_{i = 1}^{n} a_i$）

该顶点的上半部分贡献值增加（all-sum[u]），下半部分贡献值减少（sum[u]）。

#include <set>
#include <map>
#include <queue>
#include <deque>
#include <stack>
#include <cmath>
#include <cstdio>
#include <vector>
#include <string>
#include <cstring>
#include <fstream>
#include <iostream>
#include <algorithm>
using namespace std;

#define eps 1e-8
#define pb push_back
#define PI acos(-1.0)
#define INF 0x3f3f3f3f
#define clr(a,b) memset(a,b,sizeof(a)
#define bugc(_) cerr << (#_) << " = " << (_) << endl
#define FAST_IO ios::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL)

const int N=2e5+10;
typedef long long ll;
typedef unsigned long long ull;
ll a[N],sum[N],res,cnt,ans[N];

vector <int> E[N];

void dfs(int u,int fa,ll len){
    res+=len*a[u];
    sum[u]=a[u];
    for(int i=0;i<E[u].size();i++){
        int v=E[u][i];
        if(v==fa) continue;
        dfs(v,u,len+1);
        sum[u]+=sum[v];
    }
}

void DFS(int u,int fa){
    if(fa!=0) ans[u]=ans[fa]+cnt-2*sum[u];
    for(int i=0;i<E[u].size();i++){
        int v=E[u][i];
        if(v==fa) continue;
        DFS(v,u);
    }
}

int main(){
    FAST_IO;
    int n;
    cin>>n;
    for(int i=1;i<=n;i++) cin>>a[i],cnt+=a[i];
    for(int i=1;i<n;i++){
        int u,v;
        cin>>u>>v;
        E[u].push_back(v);
        E[v].push_back(u);
    }
    dfs(1,0,0);
    ans[1]=res;
    DFS(1,0);
    cout<<*max_element(ans+1,ans+1+n)<<endl;
    return 0;
}