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;
}