Codeforces 324F - Maximum White Subtree （树形dp？）

Description

思路

这题的思想有点类似求树的重心。
首先思考一下如果只要求一个点的差值。那么显然，将这个点看作根节点，然后只要dfs一下，将与它连接的所有子树中各个最大差值求个和就是答案。
在dfs的过程中把每个结点的求的值记录下来，那么对于每个结点，下面的最大差值都已经求好了，就剩下上面的值。而每个点的上面的值，都可以由它们的父节点来更新。
所以两趟dfs，一次更新下面的差值，一次更新上面的差值，答案就是上下的最大差值加上自身。

#include <bits/stdc++.h>
 
using namespace std;
const int N = 2e5 + 10;
typedef long long ll;
#define endl '\n'
#define inf 0x3f3f3f3f
const int M = 998244353;
 
int col[N];
vector<int> np[N];
int ans[N];
int down[N];
int up[N];
 
int work(int p, int fa) {
    int tot = 0;
    for(auto nt : np[p]) {
        if(nt == fa) continue;
        int v = work(nt, p);
        if(v > 0) tot += v; 
    }
    down[p] = tot;
    return tot += col[p] ? 1 : -1;
}
 
void update(int p, int fa) {
    ans[p] = max(ans[p], up[p] + down[p] + (col[p] ? 1 : -1));
    int tot = 0;
    for(auto nt : np[p]) {
        if(nt == fa) {
            if(up[p] > 0) tot += up[p];
            continue;
        }
        int num = down[nt] + (col[nt] ? 1 : -1);
        if(num > 0) tot += num;
    }
    //cout << p << ":" << tot << endl;
    for(auto nt : np[p]) {
        if(nt == fa) continue;
        int num = down[nt] + (col[nt] ? 1 : -1);
        up[nt] = tot - (num > 0 ? num : 0) + (col[p] ? 1 : -1); 
        update(nt, p);
    }
}
 
int main() {
    ios::sync_with_stdio(false);
    int n;
    cin >> n;
    for(int i = 1; i <= n; i++) {
        cin >> col[i];
        if(col[i]) ans[i] = 1;
        else ans[i] = -1;
    }
    for(int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        np[u].push_back(v);
        np[v].push_back(u);
    }
    work(1, 0);
    update(1, 0);
    for(int i = 1; i <= n; i++) cout << ans[i] << " ";
}