Codeforces Round #527 (Div. 3) F. Tree with Maximum Cost 树

 题目：

传送门

题意：

给出一颗树和每个节点的权值，求树的最大花费，花费计算方法：每个节点的高度*权值的总和，注意根节点高度为0.

思路：
两次dfs，第一次dfs求出以1为根节点的花费，并且求出每个节点的与其相连的顶点的权值和。

第二次dfs，以每个被作为根节点的顶点进行dfs，计算出每个与其相连的顶点作为根的花费。

然后遍历取最大值即可。

代码如下：

#include <bits/stdc++.h>
using namespace std;
const int maxn=2*1e5+5;
typedef long long ll;
ll a[maxn];
int n;
vector<int>v[maxn];
int vis[maxn];
ll sum[maxn];
ll ans[maxn];
ll tsum=0;
ll root1=0;
void dfs1(int now,int pre,int level)
{
    root1+=a[now]*level;
    for (int i=0;i<v[now].size();i++)
    {
        int u=v[now][i];
        if(u==pre)
            continue;
        dfs1(u,now,level+1);
        sum[now]+=sum[u];
    }
    sum[now]+=a[now];
}
void dfs2(int now,int pre)
{
    for (int i=0;i<v[now].size();i++)
    {
        int u=v[now][i];
        if(u==pre)
            continue;
        ans[u]=ans[now]-sum[u]+tsum-sum[u];
        dfs2(u,now);
    }
}
int main()
{
    memset (vis,0,sizeof(vis));
    memset (sum,0,sizeof(sum));
    scanf("%d",&n);
    for (int i=1;i<=n;i++)
    {
        scanf("%lld",&a[i]);
        tsum+=a[i];
    }
    for (int i=1;i<=n-1;i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        v[x].push_back(y);
        v[y].push_back(x);
    }
    dfs1(1,0,0);
    ans[1]=root1;
    memset (vis,0,sizeof(vis));
    dfs2(1,0);
    ll Max=-1;
    for (int i=1;i<=n;i++)
    {
        Max=max(ans[i],Max);
    }
    printf("%lld\n",Max);
    return 0;
}