[CF1324F] Maximum White Subtree

题目链接

Solution

如果只询问根节点的答案，那是一道很简单的 $dp$ 题：设 $dp[i]$ 是以 $i$ 为根的子树的答案，$dfs$ 一遍就能得到答案。

现在考虑换根的问题，如果根节点从 $x$ 换成与之相邻的 $y$，可以发现 $dp$ 值改变的只有 $dp[x]$ 和 $dp[y]$，且 dp[x] -= max(0, dp[y]), dp[y] += max(0, dp[x]) 即可。于是再 $dfs$ 一遍得到所有点的答案。

Code

#include<algorithm>
#include<cstdio>

using namespace std;

typedef long long LL;

const int N = 200005;

int n, a[N], ans[N];

struct Edge{
    int nxt, to;
}edge[N<<1];
int head[N], edgeNum;
void addEdge(int from, int to){
    edge[++edgeNum] = (Edge){head[from], to};
    head[from] = edgeNum;
}

int res[N];
int dfs(int x, int f){
    for(int i = head[x]; i; i = edge[i].nxt){
        if(edge[i].to == f) continue;
        res[x] += max(0, dfs(edge[i].to, x));
    }
    return res[x];
}

void dfsAns(int x, int f){
    ans[x] = res[x];
    for(int i = head[x]; i; i = edge[i].nxt){
        if(edge[i].to == f) continue;
        res[x] -= max(0, res[edge[i].to]);
        res[edge[i].to] += max(0, res[x]);
        dfsAns(edge[i].to, x);
        res[edge[i].to] -= max(0, res[x]);
        res[x] += max(0, res[edge[i].to]);
    }
}

int main(){
    scanf("%d", &n);
    for(int i = 1; i <= n; i++){
        scanf("%d", &a[i]);
        res[i] = a[i] == 1 ? 1 : -1;
    }
    for(int i = 1; i < n; i++){
        int u, v; scanf("%d%d", &u, &v);
        addEdge(u, v), addEdge(v, u);
    }
    dfs(1, 0);
    dfsAns(1, 0);
    for(int i = 1; i <= n; i++)
        printf("%d ", ans[i]);
    return 0;
}