Codeforces 1324 F. Maximum White Subtree (tree dp) / Comments

F. Maximum White Subtree
time limit per test2 seconds
memory limit per test256 megabytes
inputstandard input
outputstandard output
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⋅105) — 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
inputCopy
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
outputCopy
2 2 2 2 2 1 1 0 2
inputCopy
4
0 0 1 0
1 2
1 3
1 4
outputCopy
0 -1 1 -1

The meaning of problems:
Given n points, and each point of color, and the n-1 sides, even edges to ensure a tree / unrooted trees.
Each seeking the point at which the largest connected subgraph which cnt white -cnt black.

Ideas:

1. Classic tree dp, bigwigs say this is a template problem because I am more change root vegetables last night to write a collapse.
2. Personal habits, u represents the current node, v indicates a node.
3. Unrooted tree, then, of course, is to choose a point dfs, with a cnt array to record the contribution to i rooted subtree of the answer, apparently pushed by the leaf node to the root node, is when traversing the subtree dfs first and then update the status.
4. State transition is obvious: cnt [u] + = max (0, cnt [v]). // update within the sub-tree
5. dfs began to change again after the root, where I used to record the answers array dp, dp array contains two parts, one part is i sub-tree rooted at the contribution dp [i], and the other part is outside the sub-tree the face of the contribution of the answer, so we need to calculate when to change the root of the outside part of the sub-tree.
6. When traversing the subtree, if CNT [v] of cnt [u] need to contribute to cnt [u] - = cnt [v], so cnt [u] is temporarily expressed when v is the root node in the subtree contributions outside, and at this time the state transition becomes evident: cnt [v] + = max (0, cnt [u]). // update the child tree outside
7. Then remember to cnt [v] and cnt [u] restitution, because when we record the contribution of the sub-tree with an array cnt only when backtracking.
8. Details to see it ~ Code
9. If anything is unclear talking about private letter of welcome message exchange ~

Code:

#include<bits/stdc++.h>
#define ll long long
#define inf 0x3f3f3f3f
#define mod 1000000007
#define PI acos(-1)
#define fi first
#define se second
#define lowbit(x) (x&(-x))
#define mp make_pair
#define pb push_back
#define ins insert
#define si size()
#define E exp(1.0)
#define fixed cout.setf(ios::fixed)
#define fixeds(x) setprecision(x)
#pragma GCC optimize(2)
using namespace std;
inline ll read(){ll s=0,w=1;char ch=getchar();
while(ch<'0'||ch>'9'){if(ch=='-')w=-1;ch=getchar();}
while(ch>='0'&&ch<='9') s=s*10+ch-'0',ch=getchar();return s*w;}
void put1(){ puts("YES") ;}void put2(){ puts("NO") ;}void put3(){ puts("-1"); }
ll qp(ll a,ll b, ll p){ll ans = 1;while(b){if(b&1){ans = (ans*a)%p;--b;}a =
(a*a)%p;b >>= 1;}return ans%p;}ll Inv(ll x,ll p){return qp(x,p-2,p);}
ll Cal(ll n,ll m,ll p){if (m>n) return 0;ll ans = 1;for(int i = 1; i <= m; ++i)
ans=ans*Inv(i,p)%p*(n-i+1)%p;return ans%p;}

const int manx=2e5+5;

ll col[manx],vis[manx],dp[manx],cnt[manx];
vector<ll>g[manx];
ll cnt1=0,cnt2=0,n,top;

void dfs(ll u){
    vis[u]=1;
    cnt[u]=(col[u]==1?1:-1);
    for(auto v: g[u]){
        if(vis[v]) continue;
        dfs(v);
        if(cnt[v]>0) cnt[u]+=cnt[v];
    }
}
void dfs1(ll u){
    dp[u]=cnt[u]; vis[u]=1;
    for(auto v: g[u]){
        if(vis[v]) continue;
        if(cnt[v]>0) cnt[u]-=cnt[v];
        if(cnt[u]>0) cnt[v]+=cnt[u];
        dfs1(v);
        if(cnt[u]>0) cnt[v]-=cnt[u];
        if(cnt[v]>0) cnt[u]+=cnt[v];
    }
}
int main(){

    n=read();
    for(int i=1;i<=n;i++) col[i]=read(),f[i]=i;
    for(int i=1;i<n;i++){
        ll u=read(),v=read();
        g[u].pb(v),g[v].pb(u);
    }
    dfs(1);
    for(int i=1;i<=n;i++) vis[i]=0;
    dfs1(1);
    for(int i=1;i<=n;i++)
        printf("%lld ",dp[i]);
    return 0;
}