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

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

题目来源：https://ac.nowcoder.com/acm/contest/33190#question

题意

给定一棵包含$n$个节点的无根树，每个节点具有颜色0或颜色1。求其有多少连通子图，满足该子图内所有度数为1的节点颜色相同，输出这个答案数对1e9+7取模。特别地，单独的一个节点也被看作是一个合法的子图。

分析

本题为计数问题，一般来说考虑三种解法，暴力，用数学方法计算以及动态规划。考虑到数据范围，暴力不可行，数学方法也显然行不通，因此考虑动态规划。本题显然使用树形dp。先选定任何一个节点作为这棵树的根，对于$dp$数组的定义，可以用$d{p}_{i,j}$表示以第$i$个节点为根的树中，连通子图内度数为1的节点颜色都为$j$的方案数。
对于每个$d{p}_{i,j}$的计算，我们可以通过举例来推导，举例如下：

统计度数为1的节点颜色都为1的合法方案数：
由于$colo{r}_1=1$，因此三棵子树选择任意一个方案，或不选，此时的构成的连通图都是合法的。因此，总方案数为$(10+1)\times (20+1)\times (30+1).$ 同时$d{p}_{1,1}=(10+1)\times (20+1)\times (30+1).$

统计度数为1的节点颜色都为0的合法方案数：
由于$colo{r}_1=1$，因此若只选择其中一个子树的方案，构成的子树都是不合法的，反之，方案合法。因此总方案数为$(100+1)\times (200+1)\times (300+1)-(100+200+300)-1.$由于当前节点的颜色对于其父节点所取得的方案数并不会产生如上的影响，仅会由于当前节点颜色不同而减少1方案数，因此$d{p}_{1,0}=(100+1)\times (200+1)\times (300+1)-1.$

依照上方的样例，我们即可推得$d{p}_{i,j}$以及$ans$的计算公式。

参考代码

#include<bits/stdc++.h>
using namespace std;
#define ll long long
const int mod=1e9+7;
vector<int>e[300005];
ll n,x,y,dp[300005][5],col[300005],ans;
char c;
void dfs(int root,int fa){
	dp[root][0]=1;
	dp[root][1]=1;
	ll cnt=0;
	for(auto i:e[root]){
		if(i==fa)continue;
		dfs(i,root);
		dp[root][0]=(dp[root][0]*(dp[i][0]+1))%mod;
		dp[root][1]=(dp[root][1]*(dp[i][1]+1))%mod;
		cnt=(cnt+dp[i][!col[root]])%mod;
	}
	dp[root][!col[root]]--; 
	ans=((ans+dp[root][1]+dp[root][0]-cnt)%mod+mod)%mod;//防止ans出现负数
}
int main(){
	cin>>n;
	for(int i=1;i<=n;i++){
		cin>>c;
		col[i]=c-'0';
	}
	for(int i=1;i<n;i++){
		cin>>x>>y;
		e[x].push_back(y);
		e[y].push_back(x);
	}
	dfs(1,-1);
	cout<<ans;
}