Codeforces Round #627 (Div. 3) F - Maximum White Subtree (树形dp)

**F - Maximum White Subtree **

题目描述

You are given a tree consisting of n vertices. A tree is a connected undirected graph with n−1 edges. Each vertex v of this tree has a color assigned to it (av=1 if the vertex v is white and 0 if the vertex v is black).

You have to solve the following problem for each vertex v: what is the maximum difference between the number of white and the number of black vertices you can obtain if you choose some subtree of the given tree that contains the vertex v? The subtree of the tree is the connected subgraph of the given tree. More formally, if you choose the subtree that contains cntw white vertices and cntb black vertices, you have to maximize cntw−cntb.

Input

The first line of the input contains one integer n (2≤n≤2⋅10^5) — the number of vertices in the tree.

The second line of the input contains n integers a1,a2,…,an (0≤ai≤1), where ai is the color of the i-th vertex.

Each of the next n−1 lines describes an edge of the tree. Edge i is denoted by two integers ui and vi, the labels of vertices it connects (1≤ui,vi≤n,ui≠vi).

It is guaranteed that the given edges form a tree.

Output

Print n integers res1,res2,…,resn, where resi is the maximum possible difference between the number of white and black vertices in some subtree that contains the vertex i.

Examples

Input #1
9
0 1 1 1 0 0 0 0 1
1 2
1 3
3 4
3 5
2 6
4 7
6 8
5 9
Output #1
2 2 2 2 2 1 1 0 2

Input #2
4
0 0 1 0
1 2
1 3
1 4
Output #2
0 -1 1 -1

Note

The first example is shown below:

The black vertices have bold borders.

In the second example, the best subtree for vertices 2,3 and 4 are vertices 2,3 and 4 correspondingly. And the best subtree for the vertex 1 is the subtree consisting of vertices 1 and 3.

题意

有n个点n-1条边的树，问每个点所在连通块中白黑点数目最大差。

Solution

树形dp (换根).

令dp[u]为存在i节点的连通块的最大白黑节点差.对于父亲节点和儿子节点，只有儿子节点对父亲节点有贡献才向上传递，dp[u] += max(0,dp[v]);
做完一遍dp只能求出根节点的ans，因为贡献仅能向上传递，所以要求每个节点的ans，必须进行换根。

dfs2用来换根，当前根u的子节点v，记完u的答案，再把子节点v提上去做根，更新dp[u] = dp[u] - max(0,dp[v]),dp[v] = dp[v] + max(0,dp[u]),v变成根之后记录答案，依次类推下去，然后回溯（注意回溯顺序），把所有节点全部提为根一次，得到所有的ans值。

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
using namespace std;

typedef long long ll;

const int INF = 0x3f3f3f3f;
const int SZ = 200000 + 200;

struct zt
{
	int v,nxt;
}node[SZ << 1];

int color[SZ],dp[SZ],fist[SZ],ans[SZ];
int n,temp; 

inline void build(int x,int y)
{
	node[ ++ temp].v = y;
	node[temp].nxt = fist[x];
	fist[x] = temp;
}

inline void dfs1(int u,int fa)
{
	if(color[u] == 1) dp[u] = 1;
	else dp[u] = -1;
	for(int i = fist[u];i != -1;i = node[i].nxt)
	{
		int v = node[i].v;
		if(v == fa) continue;
		dfs1(v,u);
		if(dp[v] > 0) dp[u] += dp[v]; 	
	}
}

inline void dfs2(int u,int fa)
{
	ans[u] = dp[u];
	for(int i = fist[u];i != -1;i = node[i].nxt)
	{
		int t = node[i].v;
		if(t == fa) continue;
		dp[u] -= max(0,dp[t]);
		dp[t] += max(0,dp[u]);
		dfs2(t,u);
		dp[t] -= max(0,dp[u]);
		dp[u] += max(0,dp[t]);
	}
} 

int main()
{
	int a,b;
	scanf("%d",&n);
	memset(fist,-1,sizeof(fist));
	for(int i = 1;i <= n;i ++ ) scanf("%d",&color[i]);
	for(int i = 1;i < n;i ++ )
	{
		scanf("%d%d",&a,&b);
		build(a,b);
		build(b,a);
	}
	dfs1(1,-1); 
	dfs2(1,-1);
	for(int i = 1;i <= n;i ++)
	printf("%d%c",ans[i],i == n?'\n':' ');
	return 0;
}

2020.3.15