[树形dp][换根]Maximum White Subtree

题目链接

http://codeforces.com/contest/1324/problem/F

题意

对树中每一个点(v)，求包含(v)点的子树的最大(Cnt白-Cnt黑)

题解

首先以点(1)作为树根，树形(dp)求出(dp[1])，并记为答案(ans[1])。
再换做以(1)的儿子(t)作为根，更新(dp[1]=dp[1]-max(0,dp[t]),dp[t]=dp[t]+max(0,dp[1]))，并将此时的(dp[t])记为答案(ans[t])。
对(t)的儿子以此类推。

代码

#include<bits/stdc++.h>
#pragma GCC optimize(2)
using namespace std;
typedef long long ll;

int a[200005];
vector<int> edge[200005];
int dp[200005];
int ans[200005];

void dfs(int u,int fa){
  dp[u]=a[u];
  int siz=edge[u].size();
  for(int i=0;i<siz;i++){
    int v=edge[u][i];
    if(v==fa) continue;
    dfs(v,u);
    dp[u]+=max(0,dp[v]);
  }
}

void dfs2(int u,int fa){
  ans[u]=dp[u];
  int siz=edge[u].size();
  for(int i=0;i<siz;i++){
    int v=edge[u][i];
    if(v==fa) continue;
    dp[u]-=max(0,dp[v]);
    dp[v]+=max(0,dp[u]);
    dfs2(v,u);
    dp[v]-=max(0,dp[u]);
    dp[u]+=max(0,dp[v]);
  }
}

int main()
{
    int n;scanf("%d",&n);
    for(int i=1;i<=n;i++){
        scanf("%d",&a[i]);
        if(a[i]==0) a[i]=-1;
    }
    for(int i=1;i<=n-1;i++){
        int u,v;scanf("%d%d",&u,&v);
        edge[u].push_back(v);
        edge[v].push_back(u);
    }
    dfs(1,0);
    dfs2(1,0);
    for(int i=1;i<=n;i++){
        if(i!=1) printf(" ");
        printf("%d",ans[i]);
    }puts("");
    return 0;
}