Codeforces F. Maximum White Subtree (树形dp)

传送门

题意： 给定一棵 n 个节点无根树，每个点都有一个颜色，非黑即白。现在问对于每个点 i 而言，选出一个 包含i的 连通块，使其白色点与黑色点的个数差最大(注意不是绝对值)。
1 ≤ n ≤ 2e5 ，0 ≤ a[i] ≤ 1 (a[i]为0表示黑色，为1表示白色) 。

思路：

• 因为我们需使cntw-cntb尽量大，可设定每个白色节点贡献值为1，黑色节点贡献值为-1。这样就转换成两次dp求出最大子连通块的权值和，分别为自底向上和自顶向下。
• 自底向上的话可直接维护dp[ i ]，表示以点 i 为根的子树中的最大连通块。而自顶向下时主要是为了处理自底向上过程中无法包含的 fa - u 这条树链所代表的子树。
• 在第二次（dfs）自上向下走的过程中就可直接维护答案了，需要额外维护一个sum变量，用于储存 fa - u 这条链所代表的子树的贡献。显然 sum 最小为 0，代表的意义就是不选 fa - u 这条链，此时的答案便是 ans[u] = dp[u] + sum。
• ans[ u ]的正负，就决定着能否对相邻的子节点做出贡献了，如果ans[ u ]为正的话，那么减去 dp[ v ] 就是 fa - u 这条链代表的子树的权值了。
• 详细思路参考大佬博客1，同样大佬博客2也非常简洁清晰。

代码实现：

#include<bits/stdc++.h>
#define endl '\n'
#define null NULL
#define ll long long
#define int long long
#define pii pair<int, int>
#define lowbit(x) (x &(-x))
#define ls(x) x<<1
#define rs(x) (x<<1+1)
#define me(ar) memset(ar, 0, sizeof ar)
#define mem(ar,num) memset(ar, num, sizeof ar)
#define rp(i, n) for(int i = 0, i < n; i ++)
#define rep(i, a, n) for(int i = a; i <= n; i ++)
#define pre(i, n, a) for(int i = n; i >= a; i --)
#define IOS ios::sync_with_stdio(0); cin.tie(0);cout.tie(0);
const int way[4][2] = {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
using namespace std;
const int  inf = 0x7fffffff;
const double PI = acos(-1.0);
const double eps = 1e-6;
const ll   mod = 1e9 + 7;
const int  N = 2e5 + 5;

int n, a[N], dp[N], ans[N];
vector<int> g[N];

void dfs1(int u, int fa)//自底向上
{
	dp[u] = a[u];
	for(auto v : g[u]){
		if(v == fa) continue;
		dfs1(v, u);
		dp[u] += max(dp[v], 0LL);
	}
}

void dfs2(int u, int fa, int sum)//自顶向下
{
	ans[u] = dp[u] + sum;
	for(auto v : g[u]){
		if(v == fa) continue;
		dfs2(v, u, max(0LL, ans[u]-max(0LL, dp[v])));
	}
}

signed main()
{
    IOS;
	cin >> n;
	for(int i = 1; i <= n; i ++){
		cin >> a[i];
		if(!a[i]) a[i] = -1;
	}
	for(int i = 1; i < n; i ++){
		int u, v; cin >> u >> v;
		g[u].push_back(v);
		g[v].push_back(u);
	}
	dfs1(1, 0);
	dfs2(1, 0, 0);
	for(int i = 1; i <= n; i ++) cout << ans[i] << " \n"[i==n];

    return 0;
}