这道题换根时候要计算对答案的影响,就是减去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\n",ans);
}
)