F. Maximum White Subtree 树形dp*换根

大佬博客

#include <bits/stdc++.h>
using namespace std ;
const int N = 2e5 + 10 ;
int ans[N] , a[N] , dp[N] ;
vector<int> v[N] ;
int dfs(int u , int fa)
{
  dp[u] = a[u] ;
  for(auto x : v[u])
   {
     if(x == fa) continue ;
     dfs(x , u) ;
     dp[u] += max(dp[x] , 0) ;
   }
}
int dfs1(int u , int fa , int sum)
{
  ans[u] = dp[u] + sum ;
  for(auto x : v[u])
   {
     if(x == fa) continue ;
     dfs1(x , u , max(0 , ans[u] - max(0 , dp[x]))) ;
   }
}
int main()
{
  int n ;
  cin >> n ;
  for(int i = 1; i <= n ;i ++) cin >> a[i] , a[i] = (a[i] == 0 ? -1 : 1);
  for(int i = 1 , x , y ; i < n ;i ++)
   cin >> x >> y , v[x].push_back(y) , v[y].push_back(x) ;
  dfs(1 , -2) ;
  dfs1(1 , -1 , 0) ;
  for(int i = 1; i <= n ;i ++) cout << ans[i] << " " ;
  puts("") ;
  return 0 ;
}