CF1324F题解

自己入ACM的坑也有一年多了，想想自己一年前只能做到div3的D，现在已经能把F顺利做出来了。不管以后能走多远，总想留下一些东西。
题意：给定n个结点一棵树，每个结点为黑色或白色，对于每个结点，找到包含它的一棵子树，使子树的白结点与黑结点个数之差最大。
很容易想到O(n*n)的算法，就是以每个结点为根，从上到下dfs求每颗子树的能向下延伸的最大值，然后如果子树的最大值比0大，就将父亲的ans加上它。跑n遍dfs，这样显然是会TLE的。
那么该如何优化它呢？换根DP。可以明显发觉，每个父亲的ans和他的儿子的ans之间是有关系的。我们第一遍dfs求的ans，其实已经接近最大值了，因为是树性结构，对于任一结点(除根以外，因外此时根的ans就是答案)，每个儿子对它的贡献已经算出来了，只是缺少他的父亲。所以再写一个的dfs2，从根开始，如果之前的dfs将儿子的贡献算在了父亲中，那么就将父亲的贡献也加到孩子中，此时儿子的ans已是正确的，不断向下递归，最后更新完整棵树。
实现的细节可以看我的代码。

#include<iostream>
#include<cstring>
#include<algorithm>
#include<vector>
using namespace std;
int n;
int save[200010];
vector<int> bian[200010];
int l,r;
int ans[200010];
void dfs(int x,int fa)
{
	if (save[x]==1)
	{
		ans[x]=1;
	}
	else
	{ans[x]=-1;
		
	}
	for(int i=0;i<bian[x].size();i++)
	{
		if (fa==bian[x][i])
		{
			continue;
		}
		dfs(bian[x][i],x);
		if (ans[bian[x][i]]>=0)
		{
			ans[x]+=ans[bian[x][i]];
		}
	}
 } 
 
 
 
 void dfs2(int x,int fa)
 {
 	for(int i=0;i<bian[x].size();i++)
 	{
 		if (fa==bian[x][i])
 		{
 			continue;
		 }
		 if(ans[bian[x][i]]>=0&&ans[x]-ans[bian[x][i]]>=0)
		 {
		 	ans[bian[x][i]]=ans[bian[x][i]]+ans[x]-ans[bian[x][i]];
		 }
		 else if (ans[bian[x][i]]<0&&ans[x]>=0)
		 {
		 	ans[bian[x][i]]+=ans[x];
		 }
		 dfs2(bian[x][i],x);
	 }
 }
int main()
{
	
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
	{
		scanf("%d",&save[i]);
	}
	for(int i=1;i<n;i++)
	{
		scanf("%d%d",&l,&r);
		bian[l].push_back(r);
		bian[r].push_back(l);
	}
	dfs(1,0);
	dfs2(1,0);
	for(int i=1;i<=n;i++)
	{
		cout<<ans[i]<<" ";
	}
	
}