CF1324F Maximum White Subtree

题意简述

　　题目链接

　　给定一棵无根树，每个节点要么是黑点要么是白点，要求对于每个节点u，选出包含u的一个连通子图，使cnt1-cnt2最大化，其中cnt1为该连通子图内白点数，cnt2为连通子图内黑点数。

算法概述

　　记每个节点的权值w[u]为1或-1,1表示该点为白点，-1表示该点为黑点。

　　定义dp[u]表示在以u为根的子树中选一张包含u的连通子图的所有方案中cnt1-cnt2的最大值。

　　显然dp[u]=w[u]+∑max(dp[v],0)，其中v为u的子节点。这个式子的含义就是dp[u]必然是由u点以及其儿子的最大值组成的，而其所有儿子中，若某儿子的dp值小于0，则将其加上必然会使答案变得更差，故可以不选这个儿子，所以每次加之前要拿dp[v]与0做个比较。我们记vis[u]表示在计算节点u的父亲的dp值时，节点u是否有被选，1表示被选，0表示未选。

　　定义f[u]表示在整棵树中选一张包含u的连通子图的所有方案中，cnt1-cnt2的最大值。

　　考虑f[u]的组成：一部分是以其为根的子树，即dp[u]；另一部分即是在全局中将其子树挖掉后剩下的部分。这两部分互相独立，分别取最大值即可。

　　考虑第二部分的计算：记u的父节点为fa，若vis[u]=1，则说明其包含在dp[fa]中，故这一部分即为f[fa]-dp[u]；若vis[u]=0，即其不包含在dp[fa]中，则这一部分为f[fa]。最后将这一部分的值与0取max即可。

　　所以整个算法只需要两遍dfs，第一遍自底向上计算出dp数组，第二遍以f[1]（强制1为根节点）自上向下计算出f数组，即为答案。

　　时间复杂度O(n)。

参考代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int N=2e5+10;

struct Edge{
	int to,next;
}edge[N<<1];int idx;
int h[N];

void add_edge(int u,int v){edge[++idx]={v,h[u]};h[u]=idx;}

int dp[N],f[N],vis[N],w[N];
int n;

void dfs1(int p,int fa)
{
	dp[p]=w[p];
	for(int i=h[p];~i;i=edge[i].next)
	{
		int to=edge[i].to;
		if(to==fa)continue;
		dfs1(to,p);
		if(dp[to]>0)
		{
			vis[to]=1;
			dp[p]+=dp[to];
		}
	}
}

void dfs2(int p,int fa)
{
	for(int i=h[p];~i;i=edge[i].next)
	{
		int to=edge[i].to;
		if(to==fa)continue;
		int val=f[p]-(vis[to]?dp[to]:0);
		f[to]=dp[to]+(val>0?val:0);
		dfs2(to,p);
	}
}

int main()
{
	memset(h,-1,sizeof h);
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
	{
		int x;scanf("%d",&x);
		w[i]=(x?x:-1);
	}
	for(int i=1;i<=n-1;i++)
	{
		int u,v;scanf("%d%d",&u,&v);
		add_edge(u,v);
		add_edge(v,u); 
	}
	
	dfs1(1,0);
	f[1]=dp[1];
	dfs2(1,0);
	
	for(int i=1;i<=n;i++)printf("%d ",f[i]);
	return 0;
}