CF1092F Tree with Maximum Cost（换根dp）

刷了这么久终于遇到水题，十分套路的换根dp，先自底向上维护，之后自顶向下维护一遍

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=4e5+10;
int h[N],e[N],ne[N],idx;
int sz[N];
ll f[N],g[N];
ll a[N];
ll son[N],up;
int n;
void add(int a,int b){
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
ll ans;
void dfs(int u,int fa){
    sz[u]=1;
    son[u]=a[u];
    for(int i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(j==fa)
            continue;
        dfs(j,u);
        f[u]+=f[j]+son[j];
        sz[u]+=sz[j];
        son[u]+=son[j];
    }
}
void dfs1(int u,int fa){
    int i;
    for(i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(j==fa)
            continue;
        g[j]=(up-son[j])+(f[u]+g[u]-f[j]-son[j]);
        ans=max(ans,g[j]+f[j]);
        dfs1(j,u);
    }
}
int main(){
    ios::sync_with_stdio(false);
    cin>>n;
    memset(h,-1,sizeof h);
    int i;
    for(i=1;i<=n;i++){
        cin>>a[i];
        up+=a[i];
    }
    for(i=1;i<n;i++){
        int a,b;
        cin>>a>>b;
        add(a,b);
        add(b,a);
    }
    dfs(1,-1);
    dfs1(1,-1);
    ans=max(f[1],ans);
    cout<<ans<<endl;
    return 0;
}