CF1324F - Maximum White Subtree 树形DP换根

CF1324F - Maximum White Subtree

题意

$N$个点$N-1$条边的树，每个点有对应的颜色$c_i$，$c_i=1$为白色$white$，$c_i=0$为黑色$black$
对于一个点$i$，求包含点$i$的子树中最大的$cn{t}_w-cn{t}_b$
求出所有点的结果

题解

要算每个点的情况，如果都算一遍复杂度$O(N^2)$，这里考虑换根
即随便确定一个点为根，跑一遍$dp$，然后在从这个点开始转移树根，得到答案
记$f[u]$为以某个点为根时$u$的子树中连着点$u$的最大的$cn{t}_w-cn{t}_b$
转移方程：$f[u]=\sum _{f[v]>0}f[v]$
即

void dfs1(int u, int fa) {
    f[u] = c[u] ? 1 : -1;
    for (auto &v: g[u])
        if (v != fa) {
            dfs1(v, u);
            if (f[v] > 0) f[u] += f[v];
        }
}

然后考虑换根
现在从点$u$为根，变成了以点$v$为根
那么$u$为根的时候$f[u]$中如果有$f[v]$的贡献，那么需要减去
现在以点$v$为根，所以$f[u]>0$时，$f[v]$加上$f[u]$的贡献

void dfs2(int u, int fa) {
    for (auto &v: g[u])
        if (v != fa) {
            int t1 = f[u], t2 = f[v];//记录原来的值
            if (f[v] > 0) f[u] -= f[v];//u为根的时候f[u]中如果有f[v]的贡献，那么减去
            if (f[u] > 0) f[v] += f[u];//现在以v为根，u为v的子节点，加上子节点答案
            ans[v] = f[v];//更新答案
            dfs2(v, u);
            f[u] = t1, f[v] = t2;//回溯
        }
}

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAX = 2e5 + 10;

int N;
int c[MAX], f[MAX], ans[MAX];
vector<int> g[MAX];

void dfs1(int u, int fa) {
    f[u] = c[u] ? 1 : -1;
    for (auto &v: g[u])
        if (v != fa) {
            dfs1(v, u);
            if (f[v] > 0) f[u] += f[v];
        }
}

void dfs2(int u, int fa) {
    for (auto &v: g[u])
        if (v != fa) {
            int t1 = f[u], t2 = f[v];//记录原来的值
            if (f[v] > 0) f[u] -= f[v];//u为根的时候f[u]中如果有f[v]的贡献，那么减去
            if (f[u] > 0) f[v] += f[u];//现在以v为根，u为v的子节点，加上子节点答案
            ans[v] = f[v];//更新答案
            dfs2(v, u);
            f[u] = t1, f[v] = t2;//回溯
        }
}

int main() {

    scanf("%d", &N);
    for (int i = 1; i <= N; i++) scanf("%d", &c[i]);
    for (int i = 1; i < N; i++) {
        int u, v; scanf("%d%d", &u, &v);
        g[u].push_back(v); g[v].push_back(u);
    }
    dfs1(1, 0); ans[1] = f[1];
    dfs2(1, 0);
    for (int i = 1; i <= N; i++)
        printf("%d%s", ans[i], i == N ? "\n" : " ");

    return 0;
}