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