python/TangDou/Topic/HuanGenDp/CF1324F.cpp

#include <bits/stdc++.h>
using namespace std;
using VI = vector<int>;
using VVI = vector<VI>;

VI a;
VI dp;
VI ans;
VVI e;

void dfs(int x, int fa = -1) {
    dp[x] = a[x];
    for (int to : e[x]) {
        if (to == fa) continue;
        dfs(to, x);
        dp[x] += max(0, dp[to]);
    }
}

void rdfs(int x, int fa = -1) {
    ans[x] = dp[x];
    for (int to : e[x]) {
        if (to == fa) continue;
        dp[x] -= max(0, dp[to]);
        dp[to] += max(0, dp[x]);
        rdfs(to, x);
        dp[to] -= max(0, dp[x]);
        dp[x] += max(0, dp[to]);
    }
}

int main() {
    int n;
    cin >> n;
    a = dp = ans = VI(n);
    e = VVI(n);

    for (int i = 0; i < n; ++i) {
        cin >> a[i];
        if (a[i] == 0) a[i] = -1;
    }
    for (int i = 0; i < n - 1; ++i) {
        int x, y;
        cin >> x >> y;
        --x, --y;
        e[x].push_back(y);
        e[y].push_back(x);
    }
    dfs(0);
    rdfs(0);

    for (int ret : ans) cout << ret << " ";
    cout << endl;
    return 0;
}