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

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

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

题目大意

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

样例

输入：

7
1011111
1 2
1 3
2 4
3 5
2 6
4 7

输出：

28

瞎分析

本题是计数问题，计数问题有三种基本解法，分别为暴力，数学方法以及动态规划。
而由题意，暴力对于这题是肯定行不通的，数学也没有好的解法
所以，我们采用动态规划来解决这道题。
因为这题是联通图，所以可以用树形dp
方法详见代码

代码

#include<bits/stdc++.h>
using namespace std;
#define ll long long 
const int N=301000;
const long long mod=1000000007;
int n;
int col[N],dp[N][3],ans;
vector<int>f[N];
string s;

void dfs(int u,int v){//u是当前的点，v是u的父亲节点 
	dp[u][0]=dp[u][1]=1;
	//初始化，选择颜色为1或0的情况 
	ll s1=0,s0=0;
	for(int i=0;i<f[u].size();i++){
		if(f[u][i]==v)continue;//如果f[u][i]是当前节点的父亲节点
		//则跳过，防止死循环 
		dfs(f[u][i],u); //向下查询合法的子节点 
		dp[u][1]=(1ll*dp[u][1]*(dp[f[u][i]][1]+1))%mod;
		dp[u][0]=(1ll*dp[u][0]*(dp[f[u][i]][0]+1))%mod;
		//连乘当前节点的子节点方案数
		s1=(s1+dp[f[u][i]][1]+mod)%mod;
		s0=(s0+dp[f[u][i]][0]+mod)%mod;
		//统计不合法的方案数
	}	
	if(col[u]){
		dp[u][0]--;//减去本身
		ans=(ans+dp[u][1]+mod)%mod;
		//加上总方案数
		ans=(ans-s0+dp[u][0]+mod)%mod;
		//减去不合法方案数 
	}
	else{//同上
		dp[u][1]--;
		ans=(ans+dp[u][0]+mod)%mod;
		ans=(ans-s1+dp[u][1]+mod)%mod; 
	}
}

int main(){
	cin>>n>>s;//可以选择快读 
	for(int i=0;i<s.size();i++)
		col[i+1]=s[i]-'0';//将字符串转换为颜色 
		//其实不用也行，问题不是很大 
	for(int i=1;i<n;i++){
		int a,b;
		cin>>a>>b;
		f[a].push_back(b);//存边 
		f[b].push_back(a);//a与b互相连通 
	}
	dfs(1,-1);//根节点没有父亲节点 
	cout<<(ans%mod+mod)%mod;//防止ans为负 
}