python/TangDou/Topic/HuanGenDp/CF1324F.cpp

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

// 链式前向星
int e[M], h[N], idx, w[M], ne[M];
void add(int a, int b, int c = 0) {
    e[idx] = b, ne[idx] = h[a], w[idx] = c, h[a] = idx++;
}

int f1[N];
int f2[N];
int st[N];
int c[N]; // 颜色
int n;    // 节点数量

// 以1号节点为根，跑一遍dfs,填充每个节点的cnt1-cnt2的最大值
void dfs1(int u, int fa) {
    f1[u] = c[u]; // 1：白色，-1黑色，正好与 cnt1-cnt2一致
    for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v == fa) continue;
        dfs1(v, u);
        if (f1[v] > 0) {    // 如果儿子能给点钱，那我就拿着
            st[v] = 1;      // v这个儿子给了它爸爸u钱
            f1[u] += f1[v]; // 钱要累加到一起
        }
    }
}

// 换根dp
void dfs2(int u, int fa) {
    for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v == fa) continue;
        //如果 st[v]=1,说明v的贡献包含在f1[u]里面
        int val = f2[u] - (st[v] ? f1[v] : 0);
        f2[v] = f1[v] + (val > 0 ? val : 0);
        dfs2(v, u);
    }
}

int main() {
    // 初始化链式前向星
    memset(h, -1, sizeof h);
    cin >> n;

    for (int i = 1; i <= n; i++) {
        int x;
        cin >> x;
        c[i] = (x ? x : -1); // 白色c[i]=1，黑色c[i]=-1
    }

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

    dfs1(1, 0);
    f2[1] = f1[1];
    dfs2(1, 0);

    for (int i = 1; i <= n; i++) printf("%d ", f2[i]);
    return 0;
}