“蔚来杯“2022牛客暑期多校训练营5 D题: Birds in the tree

D题: Birds in the tree

原题链接:https://ac.nowcoder.com/acm/contest/33190/D

题目大意

给定包含 $n$ 个节点的树，每个节点具有颜色 $0$ 或颜色 $1$ 。求其有多少联通子图，满足度数为 $1$ 的节点颜色相同。

题解

方便起见，将题目中的树视为一颗有根树。
$s_x$ 表示节点 $x$ 的儿子， $c_x$ 表示节点 $x$ 的颜色。

给出一个定义:
若一张联通子图满足以下条件:

• 节点全部在以节点 $x$ 为根的子树内。
• 包含节点 $x$ 。
• 子图中除 $x$ 以外的度数为 $1$ 的节点的颜色均为 $c$ 。

则称该子图为节点 $x$ 的 $c$ 可并子图。
特别的，我们认为 { x } \{x\} {x} 是节点 $x$ 的 $c$ 可并子图当且仅当 $c=c_x$ 。

举例:

则 { 1 , 2 , 6 } \{1,2,6\} {1,2,6} ， { 1 , 2 , 3 , 4 , 6 } \{1,2,3,4,6\} {1,2,3,4,6} , { 1 , 2 } \{1,2\} {1,2} 均为节点 $1$ 的紫可并子图。
但 { 1 } \{1\} {1} (子树根 $1$ 颜色不为紫), { 2 , 3 , 4 } \{2,3,4\} {2,3,4} (不包含子树根 $1$ ), { 1 , 2 , 3 , 5 } \{1,2,3,5\} {1,2,3,5} (除子树根 $1$ 外，节点 $5$ 的度数为 $1$ 但颜色非紫)不是节点 $1$ 的紫可并子图。

不难发现，合法的联通子图一定是可并子图，但可并子图不一定是合法的连通子图。
当两张可并子图的根连接后，一定会形成一张合法的连通子图(同时也是新的可并子图)。

我们设 $d{p}_{x,c}$ 表示节点 $x$ 的 $c$ 可并子图数量。
转移时，对于每个儿子可以选取一个可并子图( $d{p}_{s_x,c}$ 种)或不选( $1$ 种)，根据乘法原理相乘即可，不过注意减去全部儿子都不选且子树根颜色不同的情况。
可得转移式:
$$d{p}_{x,c}=\prod _{s_x}(d{p}_{s_x,c}+1)-[c_x==c]$$

([ $f$ ]当 $f$ 为真时值为 $1$ ， $f$ 为假时值为 $0$ )

最终答案计算时，用 全部可并子图数量 减去 可并子图不合法的数量(不合法当且仅当只选择一颗子树(即子树根度数为 $1$ )且 $c_x\ne c$ (即子树根颜色不同)) 即可。
$$ans=\sum _c(\sum _{x}d{p}_{x,c}-\sum _{s_x}d{p}_{s_x,c}\ast [c!=c_x])$$

参考代码

#include<bits/stdc++.h>
using namespace std;

template<class T>inline void read(T&x){
	char c,last=' ';
	while(!isdigit(c=getchar()))last=c;
	x=c^48;
	while(isdigit(c=getchar()))x=(x<<3)+(x<<1)+(c^48);
	if(last=='-')x=-x;
}

const int MAXN=3e5+5,P=1e9+7;
int n;
int c[MAXN];
vector<int>e[MAXN],s[MAXN];
int dp[MAXN][2];

void _dfs(int u,int fa){//处理出儿子 
	for(int v:e[u]){
		if(v==fa)continue;
		s[u].push_back(v);
		_dfs(v,u);
	}
}

void dfs(int u){
	dp[u][0]=dp[u][1]=1;
	for(int v:s[u]){
		dfs(v);
		for(int C=0;C<2;++C){
			dp[u][C]=1ll*dp[u][C]*(1+dp[v][C])%P;//合并儿子的可并子图
		}
	}
	dp[u][c[u]^1]=(1ll*dp[u][c[u]^1]-1+P)%P;//减去子树根异色的情况,注意取模前+P 
}

int main()
{
	read(n);
	string str;cin>>str;str=" "+str;
	for(int i=1;i<=n;++i)c[i]=str[i]-'0';
	for(int i=1,u,v;i<n;++i){
		read(u),read(v);
		e[u].push_back(v);
		e[v].push_back(u);
	}
	_dfs(1,0);
	dfs(1);
	long long ans=0;
	for(int C=0;C<2;++C){
		for(int u=1;u<=n;++u){
			ans=(ans+dp[u][C])%P;//可并子图数
			if(c[u]!=C){//若子树根异色则非法，减去
				for(int v:s[u]){
					ans=(ans-dp[v][C]+P)%P;
				}
			}
		}
	}
	cout<<ans<<'\n';
	return 0;
}