换根dp模板

#include<bits/stdc++.h>
using namespace std;
const int MAXN=200005;
vector<int>G[MAXN];
int sz[MAXN],dp[MAXN],val[MAXN],n,u,v,ans[MAXN];
///----------------------------
void dp_link(int root1,int root2)
{
    dp[root1]+=max(dp[root2],0);
}
void dp_cut(int root1,int root2)
{
    dp[root1]-=max(dp[root2],0);
}
void get_leaf(int root)
{
     dp[root]=val[root]?1:-1;
}
///-----------------------------
void change_root(int root1,int root2)
{
    dp_cut(root1,root2);
    dp_link(root2,root1);
}
void dp_dfs(int x,int fa)
{
    get_leaf(x);
    for(auto &i:G[x])
    {
        if(i!=fa)
        {
            dp_dfs(i,x);
            dp_link(x,i);
        }
    }
    return;
}
void dfs(int x,int fa)
{
    ans[x]=dp[x];
    for(auto &i:G[x])
    {
        if(i!=fa)
        {
            change_root(x,i);
            dfs(i,x);
            change_root(i,x);
        }
    }
    return;
}
int main()
{
    scanf("%d",&n);
    for(int i=1;i<=n;++i)
    {
        scanf("%d",&val[i]);
    }
    for(int i=1;i<n;++i)
    {
        scanf("%d %d",&u,&v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dp_dfs(1,-1);
    dfs(1,-1);
    for(int i=1;i<=n;++i)
    {
        printf("%d%c",ans[i],i==n?'\n':' ');
    }
    return 0;
}