Codeforces 1324 F. Maximum White Subtree (樹形dp) /詳解

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

題意：
給定n個點，及各個點的顏色，還有n-1條邊，保證連邊成樹/無根樹。
求各個點所在的連通子圖裏面最大的cnt白-cnt黑。

思路:

1. 經典樹形dp，大佬們都說是模板題，由於我比較菜昨晚換根給寫崩了。
2. 個人習慣，u表示當前節點，v表示下一個結點。
3. 無根樹的話，首先肯定是選一個點進行dfs，用一個cnt數組來記錄以i爲根的子樹對答案的貢獻，很明顯是由葉子結點推到根節點，在遍歷子樹時是先dfs再更新狀態。
4. 狀態轉移很明顯： cnt[u]+=max(0,cnt[v]) 。//更新子樹內
5. dfs了一遍之後就開始換根，這裏我是用dp數組來記錄答案，dp數組包含兩個部分，一個部分是以i爲根的子樹對dp[i]的貢獻，另一部分是子樹外面對答案的貢獻，所以我們在換根的時候需要計算出子樹外面的部分。
6. 在遍歷子樹時， 如果cnt[v]對cnt[u]有貢獻就需要使 cnt[u]-=cnt[v]，這樣cnt[u]就暫時表示了以v爲根結點時子樹外的貢獻， 然後這時候狀態轉移就很明顯了 : cnt[v]+=max(0,cnt[u]) 。//更新子樹外
7. 然後回溯時記得把cnt[v]和cnt[u]恢復原狀，因爲我們用cnt數組僅僅時記錄子樹的貢獻 。
8. 詳細就見代碼吧 ~
9. 如果有什麼講的不清楚的歡迎留言私信交流~

代碼:

#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;
}