[树形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;
}