题意简述
给定一棵无根树,每个节点要么是黑点要么是白点,要求对于每个节点u,选出包含u的一个连通子图,使cnt1-cnt2最大化,其中cnt1为该连通子图内白点数,cnt2为连通子图内黑点数。
算法概述
记每个节点的权值w[u]为1或-1,1表示该点为白点,-1表示该点为黑点。
定义dp[u]表示在以u为根的子树中选一张包含u的连通子图的所有方案中cnt1-cnt2的最大值。
显然dp[u]=w[u]+∑max(dp[v],0),其中v为u的子节点。这个式子的含义就是dp[u]必然是由u点以及其儿子的最大值组成的,而其所有儿子中,若某儿子的dp值小于0,则将其加上必然会使答案变得更差,故可以不选这个儿子,所以每次加之前要拿dp[v]与0做个比较。我们记vis[u]表示在计算节点u的父亲的dp值时,节点u是否有被选,1表示被选,0表示未选。
定义f[u]表示在整棵树中选一张包含u的连通子图的所有方案中,cnt1-cnt2的最大值。
考虑f[u]的组成:一部分是以其为根的子树,即dp[u];另一部分即是在全局中将其子树挖掉后剩下的部分。这两部分互相独立,分别取最大值即可。
考虑第二部分的计算:记u的父节点为fa,若vis[u]=1,则说明其包含在dp[fa]中,故这一部分即为f[fa]-dp[u];若vis[u]=0,即其不包含在dp[fa]中,则这一部分为f[fa]。最后将这一部分的值与0取max即可。
所以整个算法只需要两遍dfs,第一遍自底向上计算出dp数组,第二遍以f[1](强制1为根节点)自上向下计算出f数组,即为答案。
时间复杂度O(n)。
参考代码
#include <iostream> #include <cstdio> #include <cstring> #include <algorithm> using namespace std; const int N=2e5+10; struct Edge{ int to,next; }edge[N<<1];int idx; int h[N]; void add_edge(int u,int v){edge[++idx]={v,h[u]};h[u]=idx;} int dp[N],f[N],vis[N],w[N]; int n; void dfs1(int p,int fa) { dp[p]=w[p]; for(int i=h[p];~i;i=edge[i].next) { int to=edge[i].to; if(to==fa)continue; dfs1(to,p); if(dp[to]>0) { vis[to]=1; dp[p]+=dp[to]; } } } void dfs2(int p,int fa) { for(int i=h[p];~i;i=edge[i].next) { int to=edge[i].to; if(to==fa)continue; int val=f[p]-(vis[to]?dp[to]:0); f[to]=dp[to]+(val>0?val:0); dfs2(to,p); } } int main() { memset(h,-1,sizeof h); scanf("%d",&n); for(int i=1;i<=n;i++) { int x;scanf("%d",&x); w[i]=(x?x:-1); } for(int i=1;i<=n-1;i++) { int u,v;scanf("%d%d",&u,&v); add_edge(u,v); add_edge(v,u); } dfs1(1,0); f[1]=dp[1]; dfs2(1,0); for(int i=1;i<=n;i++)printf("%d ",f[i]); return 0; }