Codeforces 1324F Maximum White Subtree

题目链接

题解：经典的树形DP，跑两遍即可，第一遍记录出每个点为根时向下所能构成的最大的值(cnt_{w}-cnt_{b})，第二遍统计出每个点为根时的答案，此时的根包括其父，所以在第二次dfs时要减去再回溯

#include<bits/stdc++.h>
using namespace std;
#define ms(x,y) memset(x, y, sizeof(x))
#define lowbit(x) ((x)&(-x))
typedef long long LL;
typedef pair<int,int> pii;


void run_case() {
    int n; cin >> n;
    vector<int> a(n), dp(n), ans(n);
    for(int i = 0; i < n; ++i) {
        cin >> a[i];
        if(a[i] == 0) a[i] = -1;
    }
    vector<vector<int>> G(n+1);
    for(int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        u--, v--;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    function<void(int, int)> dfs1 = [&](int u, int fa) {
        dp[u] = a[u];
        for(auto v: G[u]) {
            if(v == fa) continue;
            dfs1(v, u);
            dp[u] += max(0, dp[v]);
        }
    };
    function<void(int, int)> dfs2 = [&](int u, int fa) {
        ans[u] = dp[u];
        for(auto v: G[u]) {
            if(v == fa) continue;
            dp[u] -= max(0, dp[v]);
            dp[v] += max(0, dp[u]);
            dfs2(v, u);
            dp[v] -= max(0, dp[u]);
            dp[u] += max(0, dp[v]);
        }
    };
    dfs1(1, -1);
    dfs2(1, -1);
    for(auto i : ans) cout << i << " ";
}

int main() {
    ios::sync_with_stdio(false), cin.tie(0);
    cout.flags(ios::fixed);cout.precision(2);
    //int t; cin >> t;
    //while(t--)
    run_case();
    cout.flush();
    return 0;
}