题目描述

难度分:$1800$

输入$n(2 \leq n \leq 2 \times 10^5)$和长为$n$的数组$a$，只包含$0$和$1$。

然后输入一棵无向树的$n-1$条边，节点编号从$1$到$n$。

$a[i]$表示节点$i$的颜色，$0$表示黑色，$1$表示白色。

定义$f(i)$为包含节点$i$的连通块中，白色节点个数减去黑色节点个数的最大值。输出$f(1),f(2),…,f(n)$，注意$f(i)$可能是负数，见样例$2$。

输入样例$1$

9
0 1 1 1 0 0 0 0 1
1 2
1 3
3 4
3 5
2 6
4 7
6 8
5 9

输出样例$1$

2 2 2 2 2 1 1 0 2 

输入样例$2$

4
0 0 1 0
1 2
1 3
1 4

输出样例$2$

0 -1 1 -1 

算法

换根DP

比较明显的一个换根DP题目，因为每个节点都要求个最大值。

树形DP

状态定义

$f[u]$表示仅考虑以$u$为根的子树，在包括节点$u$的情况下，能找到的连通块“白-黑”最大值。

状态转移

如果$u$是白色，自增$f[u]$，否则自减$f[u]$。然后遍历$u$的所有子节点$v$，如果$f[v] \gt 0$就将其累加在$f[u]$上。状态转移方程为$f[u]=delta+\Sigma_{v}max(0,f[v])$，其中$a[u]=1$时$delta=1$，$a[u]=0$时$delta=-1$。

DFS换根

换根函数为$dfs(u,vx)$，其中$vx$表示$u$上面的节点（不包括$u$）能够贡献的“白-黑”最大值。还是遍历$u$的所有子节点$v_i$，对于一个子节点$v_i$，需要从$u$开始往$vx$考虑，将$vx$作为根，那么$vx$对于$v_i$节点就要更新为$vx+\Sigma_{j=0}^{i-1}max(0,f[v_{j}])+\Sigma_{j=i+1}^{cnt-1}max(0,f[j])+delta$。

其中$a[u]=1$时$delta=1$，$a[u]=0$时$delta=-1$，$cnt$是$u$的子节点数目。此时如果$vx \lt 0$，对于$v_i$来说那还不如舍弃$v_i$上面的贡献，所以它的$vx$应该是$max(0,vx+\Sigma_{j=0}^{i-1}max(0,f[v_{j}])+\Sigma_{j=i+1}^{cnt-1}max(0,f[j])+delta)$。这两段求和操作可以通过前后缀分解先预处理出来，从而$O(1)$计算这两段求和。

节点$u$的最终答案就应该是$g[u]=f[u]+vx$。

复杂度分析

时间复杂度

本质上就是对整棵无向树进行了两次DFS，因此算法的时间复杂度为$O(n)$。

空间复杂度

树的邻接表空间复杂度为$O(n)$；$f$和$g$两个DP数组空间复杂度为$O(n)$；在DFS的过程中需要前后缀分解，需要开辟的前后缀数组总长度在$O(n)$级别；DFS在最差情况下递归深度在$O(n)$级别（树退化成一条链）。因此，整个算法的额外空间复杂度为$O(n)$。

C++ 代码

#include <bits/stdc++.h>

using namespace std;
const int N = 200010;
int n, a[N], f[N], g[N];
vector<int> graph[N];

void dfs1(int u, int fa) {
    f[u] = a[u]? 1: -1;
    for(int v: graph[u]) {
        if(v == fa) continue;
        dfs1(v, u);
        if(f[v] > 0) f[u] += f[v];
    }
}

void dfs2(int u, int fa, int vx) {
    g[u] = f[u] + vx;
    vector<int> children;
    for(int v: graph[u]) {
        if(v == fa) continue;
        children.push_back(f[v]);
    }
    int sz = children.size();
    vector<int> pre(sz), suf(sz);
    for(int i = 1, j = sz - 2; i < sz; i++, j--) {
        pre[i] = pre[i - 1] + max(0, children[i - 1]);
        suf[j] = suf[j + 1] + max(0, children[j + 1]);
    }
    int i = 0;
    for(int v: graph[u]) {
        if(v == fa) continue;
        int nxt = vx + (a[u]? 1: -1) + pre[i] + suf[i];
        dfs2(v, u, max(0, nxt));
        i++;
    }
}

int main() {
    scanf("%d", &n);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &a[i]);
        graph[i].clear();
        f[i] = 0;    // f[i]表示以i为根的子树中，在包含i的情况下，连通块中“白-黑”的最大值
    }
    for(int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
    dfs1(1, 0);
    dfs2(1, 0, 0);
    for(int i = 1; i <= n; i++) {
        printf("%d ", g[i]);
    }
    return 0;
}