python/TangDou/Topic/HuanGenDp/CF1324F.cpp

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int N = 2e5 + 10;

struct Edge {
    int to, next;
} edge[N << 1];
int idx;
int h[N];

void add_edge(int u, int v) {
    edge[++idx] = {v, h[u]};
    h[u] = idx;
}

int dp[N], f[N], vis[N], w[N];
int n;

void dfs1(int p, int fa) {
    dp[p] = w[p];
    for (int i = h[p]; ~i; i = edge[i].next) {
        int to = edge[i].to;
        if (to == fa) continue;
        dfs1(to, p);
        if (dp[to] > 0) {
            vis[to] = 1;
            dp[p] += dp[to];
        }
    }
}

void dfs2(int p, int fa) {
    for (int i = h[p]; ~i; i = edge[i].next) {
        int to = edge[i].to;
        if (to == fa) continue;
        int val = f[p] - (vis[to] ? dp[to] : 0);
        f[to] = dp[to] + (val > 0 ? val : 0);
        dfs2(to, p);
    }
}

int main() {
    memset(h, -1, sizeof h);
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) {
        int x;
        scanf("%d", &x);
        w[i] = (x ? x : -1);
    }
    for (int i = 1; i <= n - 1; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        add_edge(u, v);
        add_edge(v, u);
    }

    dfs1(1, 0);
    f[1] = dp[1];
    dfs2(1, 0);

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