【題解】Codeforces 1324F. Maximum White Subtree

#include <bits/stdc++.h>
using namespace std;
 
int n, a[200005], dp[200005], ans[200005];
vector <int> g[200005];
 
int dfs(int x, int pre){
    if (a[x]) dp[x] = 1;
    else dp[x] = -1;
    for (int i:g[x]){
        if (i == pre) continue;
        int ret = dfs(i, x);
        if (ret > 0) dp[x] += ret;
    }
    return dp[x];
}
void dfs2(int x, int pre){
    for (int i:g[x]){
        if (i == pre) continue;
        int tmp = ans[x];
        if (dp[i] > 0) tmp -= dp[i];
        ans[i] = dp[i];
        if (tmp > 0) ans[i] += tmp;
        dfs2(i, x);
    }
}
 
int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> n;
    for (int i = 1; i <= n; i++){
        cin >> a[i];
    }
    int x, y;
    for (int i = 1; i < n; i++){
        cin >> x >> y;
        g[x].push_back(y);
        g[y].push_back(x);
    }
    dfs(1, 0);
    ans[1] = dp[1];
    dfs2(1, 0);
    cout << ans[1];
    for (int i = 2; i <= n; i++){
        cout << " " << ans[i];
    }
    cout << "\n";
}