python/TangDou/Topic/HuanGenDp/CF1324F.cpp

#include <bits/stdc++.h>
using namespace std;
const int N = 200010, 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 n, c[N];

// 设白点个数cnt1,黑点个数cnt2, f[u]：以u为根的子树中，cnt1-cnt2的最大值
int f1[N], f2[N];

void dfs1(int u, int fa) {
    if (c[u])
        f1[u] = 1; // 如果u是白色则cnt1=1,那么目前黑色数量cnt2=0,cnt1-cnt2是固定值=1
    else
        f1[u] = -1; // 如果u是黑色则cnt1=0,那么目前黑色数量cnt2=1,cnt1-cnt2是固定值=0-1=-1

    for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v == fa) continue;
        dfs1(v, u);
        f1[u] += max(f1[v], 0); // 为了贪图cnt1-cnt2的最大值，那么对结果有贡献的就竞争一下，如果都小于0了，就退出评比
    }
}

// 换根dp
void dfs2(int u, int fa) {
    for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v == fa) continue;
        // 本题的核心:递推式
        f2[v] = max(f2[u] + f1[u] - max(f1[v], 0), 0);
        dfs2(v, u);
    }
}
int main() {
    // 初始化链式前向星
    memset(h, -1, sizeof h);
    cin >> n;

    for (int i = 1; i <= n; i++) cin >> c[i]; // 每个节点都有自己的颜色c[i]

    for (int i = 1; i < n; i++) { // n-1条边
        int a, b;
        cin >> a >> b;
        add(a, b), add(b, a);
    }
    // 第一次dfs
    dfs1(1, 0);
    // 第二次dfs,换根
    dfs2(1, 0);
    // 输出
    for (int i = 1; i <= n; i++) printf("%d ", f1[i] + f2[i]);
    return 0;
}