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

题目链接

题意

给定一棵树，树上的n个节点有黑白两色，需要你计算出各节点的子序列中白点总数减黑点总数的最大值。

思路

其实很容易想到以某一结点为根节点，当其子节点的值可以使其更大时，则进行更新，否则保持不变。用DP维护该点下方的子序列的最大值即可获取答案的一部分，即除根节点外，其他节点以下部分子序列的最大值。此时若要获得所有节点的答案，则要进行自上而下的维护。再由之前的根节点出发，若其答案减去当前遍历到的子节点的贡献依然大于零，则其可以对子节点上方序列的答案进行补充，依次向下直到更新完毕即可获取答案。

代码

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int mod=1e9+5;
const int maxn=1e6+5;
vector<int>ve[maxn];
int n,a[maxn],dp[maxn],ans[maxn];
void dfs(int x,int fa)
{
    dp[x]=a[x];
    for(int i=0;i<ve[x].size();i++)
    {
        int to=ve[x][i];
        if(to==fa)
            continue;
        dfs(to,x);
        dp[x]+=max(0,dp[to]);
    }
}
void dfs1(int x,int fa)
{
    for(int i=0;i<ve[x].size();i++)
    {
        int to=ve[x][i];
        if(to==fa)
            continue;
        ans[to]=dp[to]+max(0,ans[x]-max(0,dp[to]));
        dfs1(to,x);
    }
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    cin>>n;
    for(int i=1;i<=n;i++)
    {
        cin>>a[i];
        if(!a[i])
            a[i]=-1;
    }
    for(int i=1;i<n;i++)
    {
        int u,v;
        cin>>u>>v;
        ve[u].push_back(v);
        ve[v].push_back(u);
    }
    dfs(1,-1);
    ans[1]=dp[1];
    dfs1(1,-1);
    for(int i=1;i<=n;i++)
        cout<<ans[i]<<" ";
    puts("");
    return 0;
}