CF1324F - Maximum White Subtree 树形DP换根

CF1324F - Maximum White Subtree

题意

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

题解

要算每个点的情况，如果都算一遍复杂度 $O(N^2)$，这里考虑换根
即随便确定一个点为根，跑一遍 $dp$，然后在从这个点开始转移树根，得到答案
记 $f[u]$为以某个点为根时 $u$的子树中连着点 $u$的最大的 $cnt_w-cnt_b$
转移方程： $f[u]=\displaystyle\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;
}