Codeforces Round #527 (Div. 3) F. Tree with Maximum Cost（树形换根DP）

题目链接

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=2e5+5;
vector<int>g[maxn];
ll cnt=0,a[maxn],dp[maxn],sum[maxn];
void dfs(int x,int fa,int deep)
{
	dp[x]=deep*a[x];sum[x]=a[x];
	for(int to:g[x])
	{
		if(to==fa) continue;
		dfs(to,x,deep+1);
		sum[x]+=sum[to];
		dp[x]+=dp[to];
	}
}
void dfs2(int x,int fa)
{
	for(int to:g[x])
	{
		if(to==fa) continue;
		dp[to]+=dp[x]-dp[to]+(cnt-sum[to])-sum[to];
		dfs2(to,x);
	}
}
int main()
{
	int n,u,v;
	scanf("%d",&n);
	for(int i=1;i<=n;++i) scanf("%lld",&a[i]),cnt+=a[i];
	for(int i=1;i<n;++i)
	{
		scanf("%d %d",&u,&v);
		g[u].push_back(v);
		g[v].push_back(u);
	}
	dfs(1,-1,0);
	dfs2(1,-1);
	ll ans=0;
	for(int i=1;i<=n;++i) ans=max(ans,dp[i]);
	printf("%lld\n",ans);
}