python/TangDou/Topic/HuanGenDp/CF1324F.cpp

#include <bits/stdc++.h>
using namespace std;
const int N = 200010, M = N << 1;

int n, a[N], f1[N], f2[N];
int head[N << 1], nxt[N << 1], to[N << 1], tot = 0;
void add(int u, int v) {
    to[++tot] = v;
    nxt[tot] = head[u];
    head[u] = tot;
}

void dfs1(int u, int fa) {
    if (a[u])
        f1[u] = 1;
    else
        f1[u] = -1;

    for (int i = head[u]; i; i = nxt[i]) {
        int v = to[i];
        if (v == fa)
            continue;
        dfs1(v, u);
        f1[u] += max(f1[v], 0);
    }
}

void dfs2(int u, int fa) {
    for (int i = head[u]; i; i = nxt[i]) {
        int v = to[i];
        if (v == fa)
            continue;
        f2[v] = max(f2[u] + f1[u] - max(f1[v], 0), 0);
        dfs2(v, u);
    }
}
int main() {
    cin >> n;

    for (int i = 1; i <= n; i++) cin >> a[i];

    for (int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;
        add(a, b), add(b, a);
    }

    dfs1(1, 1);

    dfs2(1, 1);
    for (int i = 1; i <= n; i++) printf("%d ", f1[i] + f2[i]);

    return 0;
}