cf codeforces round#527F. Tree with Maximum Cost树形dp

这道题换根时候要计算对答案的影响，就是减去to节点的子树和sum，加上from节点的子树和sum（这里假设to和sum都是一棵树的根，这棵树没有其他部分

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
int n;
ll a[200010];
ll sum[200010];
ll ans,res;
vector<int> vec[200010];
void dfs1(int v,int p,int h)
{
	sum[v]=a[v];
	res+=a[v]*h;
	for(int i=0;i<vec[v].size();i++)
	{
	    int to=vec[v][i];
	    if(to==p)
	       continue;
	    dfs1(to,v,h+1);
	    sum[v]+=sum[to];
	}
}
void dfs2(int v,int p)
{
	ans=max(ans,res);
	for(int i=0;i<vec[v].size();i++)
	{
		int to=vec[v][i];
		if(to==p)
		  continue;
		res-=sum[to];
		sum[v]-=sum[to];
		res+=sum[v];
		sum[to]+=sum[v];
		dfs2(to,v);
		sum[to]-=sum[v];
		res-=sum[v];
		sum[v]+=sum[to];
		res+=sum[to];
	}
}
int main()
{
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
	{
		scanf("%lld",&a[i]);
	}
	for(int i=1;i<=n-1;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		vec[u].push_back(v);
		vec[v].push_back(u);
	}
	dfs1(1,-1,0);
	dfs2(1,-1);
	printf("%lld
",ans);
}

　　

）