换根dp

Posted 2021-03-05 reshuffle

https://codeforces.com/contest/1324/problem/F

description

you are given a tree, in which the vertices are all printed either black or white. , find out the maximum difference between the number of black vertices and the number of white vertices you can obtain if you choose some subtree contain vertex v. you need to print that maximum of every vertices in the tree.

analysis

something we need to notice.

• the root are not fixed, which means we have to find every subtree based on any root.
• we are to count the number of a particular color, but we have to choose wisely.

to sum up, what we need is DP with rerooting technique.

template

the process are as follows

1. fix an arbitary root and calculate DP by simgle DFS.
2. change the root from let‘s say (u) to any vertex (v) adjacent to it. update or recalculate.

#include <bits/stdc++.h>
#pragma GCC diagnostic error "-std=c++11"
#define SIS std::ios::sync_with_stdio(false),cin.tie(0),cout.tie(0)
//#define endl '
'
#define all(A) (A).begin(),(A).end()
#define FOR(I, A, B) for (int I = (A); I <= (B); ++I)
#define PER(I, A, B) for (int I = (A); I >= (B); --I)
#define lson k*2
#define rson k*2+1
#define fi first
#define se second
#define DB(A) cout<<(A)<<endl
#define DB1(A,B,C) cout<<(A)<<" "<<(B)<<" "<<(C)<<"!"<<endl
#define PB push_back
#define Pair pair<int,int>
#define MP make_pair
#define LL long long
using namespace std;
const int maxn=3e5+10;
const int MAX=1000;
const int inf=0x3f3f3f3f;   
const int mod=1e9+7;
//head
int dp[maxn];
int note[maxn];
int ans[maxn];
vector<int>a[maxn];
void dfs(int now,int fa=0)
{
    dp[now]=note[now];
    for (auto to:a[now])
    {
        if (to==fa) continue;
        dfs(to,now);
        dp[now]+=max(0,dp[to]);
    }
}
void dfs2(int now,int fa=0)
{
    ans[now]=dp[now];//
    for (auto to:a[now])
    {
        if (to==fa) continue;
        dp[now]-=max(0,dp[to]);
        dp[to]+=max(0,dp[now]);
        dfs2(to,now);
        dp[to]-=max(0,dp[now]);
        dp[now]+=max(0,dp[to]);     
    }   
}
int main()
{
    SIS;
    int n;
    cin>>n;
    FOR(i,1,n)
    {
        int x;
        cin>>x;
        if (x==0) x=-1;
        note[i]=x;
    }
    FOR(i,1,n-1)
    {
        int x,y;
        cin>>x>>y;
        a[x].PB(y);
        a[y].PB(x);
    }
    dfs(1);
    dfs2(1);
    FOR(i,1,n) cout<<ans[i]<<" ";cout<<endl;
}