65 lines
1.6 KiB
C++
65 lines
1.6 KiB
C++
#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 f[N];
|
||
int g[N];
|
||
int c[N]; // 颜色
|
||
int n; // 节点数量
|
||
|
||
// 以1号节点为根，跑一遍dfs,填充每个节点的cnt1-cnt2的最大值
|
||
void dfs1(int u, int fa) {
|
||
f[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);
|
||
f[u] += max(0, f[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;
|
||
int val = g[u] - max(f[v], 0);
|
||
g[v] = f[v] + max(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);
|
||
}
|
||
|
||
// 第一次dfs
|
||
dfs1(1, 0);
|
||
|
||
// 它们两个是一个意思
|
||
g[1] = f[1];
|
||
// 换根dp
|
||
dfs2(1, 0);
|
||
// 输出答案
|
||
for (int i = 1; i <= n; i++) printf("%d ", g[i]);
|
||
return 0;
|
||
} |