CF1324F Maximum White Subtree 题解

题目链接：CF1324F Maximum White Subtree

题意：

给定一棵 $n$ 个节点无根树，每个节点 $u$ 有一个颜色 $a_u$，若 $a_u$ 为 $0$ 则 $u$ 是黑点，若 $a_u$ 为 $1$ 则 $u$ 是白点。

对于每个节点 $u$，选出一个包含 $u$ 的连通子图，设子图中白点个数为 $\mathtt{cnt_1}$，黑点个数为 $\mathtt{cnt_2}$，请最大化 $\mathtt{cnt_1 - cnt_2}$。并输出这个值。

$1 \leq n \leq 2 \times 10^5$，$0 \leq a_u \leq 1$。

换根dp简单题，先钦定 $1$ 为根，然后令 $a_u = 2a_u - 1$ （即 $1$ 和 $-1$ ）

设 $f_i$ 表示 $i$ 所在子树的最大值，$g_i$ 为 $i$ 为根时的最大值，则

$$f_u = a_i + \sum_{v \in \mathtt{son}(u)}f_v$$

开始换根。首先有 $g_1 = f_1$ 。

$$g_u = \max\{0, g_{\mathtt{fa}} - (0 \uparrow f_u)\} + f_u$$

时间复杂度 $\mathcal{O}(n)$

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
void up(int &x,int y) { x < y ? x = y : 0; }
void down(int &x,int y) { x > y ? x = y : 0; }
#define N ((int)(2e5+15))

int n,pos=1,head[N],col[N],f[N],g[N];
struct Edge { int u,v,next; } e[N * 2];
void addEdge(int u,int v) { e[++pos] = {u,v,head[u]}; head[u] = pos; }
void dfs1(int u,int fa)
{
    f[u] = col[u];
    for(int i=head[u]; i; i=e[i].next)
        if(e[i].v != fa) { dfs1(e[i].v,u); f[u] += max(0ll, f[e[i].v]); }
}
void dfs2(int u,int fa)
{
    if(u != 1) {g[u] = max(0ll, g[fa] - max(0ll,f[u])) + f[u]; }
    for(int i=head[u]; i; i=e[i].next) if(e[i].v != fa) dfs2(e[i].v,u);
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> n;
    for(int i=1; i<=n; i++){ cin >> col[i]; col[i] = 2 * col[i] - 1; }
    for(int i=1,u,v; i<n; i++) { cin >> u >> v; addEdge(u,v); addEdge(v,u); }
    dfs1(1,0); g[1] = f[1]; dfs2(1,0);
    for(int i=1; i<=n; i++) cout << g[i] << " \n"[i==n];
    return 0;
}