CF1092F Tree with Maximum Cost（换根dp）

传送门

---

解题思路

先预处理出以1为根节点的各节点的size。
定义一个节点的size为以这个节点为根的子树的a的和。
然后求出ans[1]，再进行换根dp求出每个节点的ans，比较大小输出即可。

AC代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#include<iomanip>
#include<vector>
#include<ctime>
#include<set>
#include<queue>
#include<stack>
#include<algorithm>
using namespace std;
template<class T>inline void read(T &x)
{
    x=0;register char c=getchar();register bool f=0;
    while(!isdigit(c))f^=c=='-',c=getchar();
    while(isdigit(c))x=(x<<3)+(x<<1)+(c^48),c=getchar();
    if(f)x=-x;
}
template<class T>inline void print(T x)
{
    if(x<0)putchar('-'),x=-x;
    if(x>9)print(x/10);
    putchar('0'+x%10);
}
const int maxn=2e5+5;
int n,cnt,f[maxn],p[maxn];
long long dp[maxn],ans,a[maxn],siz[maxn],dep[maxn];
struct node{
	int v,next;
}e[maxn*2];
void insert(int u,int v){
	cnt++;
	e[cnt].v=v;
	e[cnt].next=p[u];
	p[u]=cnt;
}
void dfs(int u,int fa){
	f[u]=fa;
	siz[u]+=a[u];
	for(int i=p[u];i!=-1;i=e[i].next){
		int v=e[i].v;
		if(v==fa) continue;
		dep[v]=dep[u]+1;
		dp[1]+=1ll*dep[v]*a[v];
		dfs(v,u);
		siz[u]+=siz[v];
	}
}
void dfs2(int u){
	ans=max(ans,dp[u]);
	for(int i=p[u];i!=-1;i=e[i].next){
		int v=e[i].v;
		if(v==f[u]) continue;
		dp[v]=dp[u]-siz[v]+siz[1]-siz[v];
		dfs2(v);
	}
}
int main(){
	//freopen(".in","r",stdin);
	//freopen(".out","w",stdout);
	ios::sync_with_stdio(false);
	memset(p,-1,sizeof(p));
	cin>>n;
	for(int i=1;i<=n;i++) cin>>a[i];
	for(int i=1;i<n;i++){
		int u,v;
		cin>>u>>v;
		insert(u,v);
		insert(v,u);
	}
	dfs(1,0);
	dfs2(1);
	cout<<ans;
	return 0;
}