[树形dp]Birds in the tree 2022牛客多校第5场 D

题目描述

One day, when NIO was walking under the tree, he found that there are birds standing on the tree, and each bird is standing on one leaf. He wondered about in how many sub-trees form by connected vertex-induced subgraphs from the original tree, all birds are in the same gender. The number would be very large, just mod 109+710^9+7109+7 in the output.

输入描述:

There are multiple input cases.

For each case, the first line contains an integer,  nnn, denoting the number of nodes on the tree. The second line is a binary string, where 111 denotes the male birds and 000 denotes the female birds.

Following by n−1n-1n−1 lines. In each line there are two integers, xi,yix_i, y_ixi​,yi​, denoting there is a path connecting node xxx and node yyy.

1≤n≤3×105,1≤xi,yi≤N1\leq n \leq 3\times10^5, 1\leq x_i,y_i \leq N1≤n≤3×105,1≤xi​,yi​≤N

输出描述:

Output a number module 109+710^9+7109+7, the number of subgraphs of the given tree that form trees where all its leaves are the same color.

示例1

输入

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

输出

28

题意： 给出一颗包含n点的树，树上每个点都有一只鸟，鸟的性别分为雄性和雌性，用一个字符串s表示这n只鸟的性别，问有多少种子图能使子图中度为1的点上的鸟同性别。

分析： 可以在树上做dp，也就是树形dp，首先设状态dp[i][j]表示以第i个点为根的子树中且必须包含第i个点的子图数，dp[i][0]表示度为1的点性别都是0的子图数，dp[i][1]表示度为1的点性别都是1的子图数。之后模拟一下样例可以得到状态转移方程，设now为当前点，to[i]为各子结点，那么dp[now][0] = (dp[to[1]][0]+1)*(dp[to[2]][0]+1)*......这是因为每个子结点都会提供dp[to[i]][0]个合法子图，再加上一个不选的情况，也就是提供一个空图，所以每个子结点会提供dp[to[i]][0]+1个方案，而不同子结点方案数是可以乘起来的，所以这个乘积就是dp[now][0]，不过这样说其实不太严谨，这其实只是方案的总数，实际上有一些方案是不符合状态含义的，我们要给它减去。有两种情况是这样的，第一种是now点性别为1，那么每个子结点全不选的情况就表示子图中只包含now自己一个点，显然一个1是不符合状态含义的，所以这时候方案数要减1。第二种不合法的情况是now点性别为1，并且只选了某一个子结点，其余子结点都不选，那么这时候子图中now点也成了度为1的点，而其余度为1的点性别都是0，也是不符状态含义的，可以发现如果now点选了两个及以上的子结点，此时now点就不是度为1的点了，所以它是什么性别都无所谓，dp初始化初始状态也就是叶子结点，它们如果性别为1那么dp[i][1] = 1否则dp[i][0] = 1，其余状态置0。

在减的时候要注意，应该先剪掉第一种情况，然后父节点用该值更新其自身dp值，然后再剪掉第二种情况，这是因为虽然剪掉两种情况才是真正符合状态含义的dp值，但是其父节点并不是用子结点真正符合状态含义的dp值来更新的，父节点是用减去第一种情况的伪dp值来更新自己dp值的，样例中就包含这种情况，如果还不明白可以动手模拟一遍样例的过程，附一张样例的dp值详细递推过程：

 

具体代码如下：

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cstring>
#include <cmath>
#include <string>
#include <vector>
#define int long long
using namespace std;

const int mod = 1e9+7;
int dp[300005][2], ans;
char s[300005];
vector<int> tr[300005];

void dfs(int now, int fa){
	if(tr[now].size() == 1){
		dp[now][s[now]-'0'] = 1;
		return;
	}
	int t0 = 1, t1 = 1;
	for(int i = 0; i < tr[now].size(); i++){
		int to = tr[now][i];
		if(to == fa) continue;
		dfs(to, now);
		t0 = (t0*(dp[to][0]+1))%mod;
		t1 = (t1*(dp[to][1]+1))%mod;
	}
	dp[now][0] = t0;
	dp[now][1] = t1;
	if(s[now] == '0')
		dp[now][1] = ((dp[now][1]-1)%mod+mod)%mod;
	else
		dp[now][0] = ((dp[now][0]-1)%mod+mod)%mod;
}

void dfs2(int now, int fa){
	int s0 = 0, s1 = 0;
	for(int i = 0; i < tr[now].size(); i++){
		int to = tr[now][i];
		if(to == fa) continue;
		s0 = (s0+dp[to][0])%mod;
		s1 = (s1+dp[to][1])%mod;
		dfs2(to, now);
	}
	ans = (ans+dp[now][0]+dp[now][1])%mod;
	if(s[now] == '0')
		ans = ((ans-s1)%mod+mod)%mod;
	else
		ans = ((ans-s0)%mod+mod)%mod;	
}

signed main()
{
	int n;
	while(~scanf("%lld", &n)){
		scanf("%s", s+1);
		for(int i = 1; i <= n; i++)
			dp[i][0] = dp[i][1] = 0;
		for(int i = 1; i < n; i++){
			int u, v;
			scanf("%lld%lld", &u, &v);
			tr[u].push_back(v);
			tr[v].push_back(u);
		}
		ans = 0;
		dfs(1, 0);
		dfs2(1, 0);
		printf("%lld\n", ans);
	}
    return 0;
}