【NOIP模拟】cut

Description

F大爷热爱切树。今天他找到一棵黑白树，不到两秒钟，F大爷就把这棵树切掉了。已知原先树上共n个点，每个点都是黑点或者白点，F大爷切去若干条边后，分成的若干个连通子树中每块恰有一个黑点， 请问有多少种切树的方案满足这个条件？两种方案不同当且仅当存在一条边在其中一个方案中被切而在另一个方案中没被切。

n<=10^5，至少有一个黑点。

Input Format

第一行一个正整数n，表示树的点数，点编号1~n

第二行n 个数，每个数字为0或1，依次表示每个点的颜色，若第i个数字为1，则i号点为黑点，否则i号点为白点。

接下来n-1行， 每行两个正整数， 表示树上的一条边连接的两点。

Output Format

输出一个整数，表示答案对10^9+7 取模后的结果。

Solution

显然树形DP，思路就不讲了，

Code

#include <cstdio>
#include <algorithm>
#include <cstring>
#define ll long long
#define N 100010
using namespace std;

const int yh=1e9+7;
struct info{int to,nex;}e[N*2];
int n,tot,head[N*2],c[N],in[N],cnt;
ll f[N][2],Ans;

inline void Link(int u,int v){
	e[++tot].to=v;
	e[tot].nex=head[u];
	head[u]=tot;
}

inline int read() {
	int x=0,f=1;char ch=getchar();
	while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
	return x*f;
}

ll DP(int u,int x,int fa){
	ll &tmp=f[u][x];
	if(tmp!=-1) return tmp;
	tmp=!x;

	for(int i=head[u];i;i=e[i].nex){
		int v=e[i].to;
		if(v==fa) continue;

		if(x==0){
			if(c[v]==1) tmp=(tmp*1ll*DP(v,0,u))%yh;
			else tmp=(tmp*1ll*(DP(v,0,u)+DP(v,1,u)))%yh;
		}else{
			ll sum=1;
			if(c[v]==1) sum=(sum*1ll*DP(v,0,u))%yh;
			else sum=(sum*1ll*DP(v,1,u))%yh;	

			for(int j=head[u];j;j=e[j].nex){
				int tv=e[j].to;
				if(tv==v||tv==fa) continue;
				if(c[tv]==1) sum=(sum*1ll*DP(tv,0,u))%yh;
				else sum=(sum*1ll*(DP(tv,0,u)+DP(tv,1,u)))%yh;
			}
			tmp=(tmp+sum)%yh;
		}
	}
	return tmp%yh;
}

int main(){
	n=read();
	for(int i=1;i<=n;++i) c[i]=read(),cnt+=c[i];
	for(int i=1;i<n;++i){
		int u=read(),v=read();
		in[u]++,in[v]++;
		Link(u,v);
		Link(v,u);
	}

	if(in[1]==n-1){
		printf("%d\n",(c[1])?1:cnt);
		return 0;
	}

	memset(f,-1,sizeof(f));
	Ans=(c[1]==1)?DP(1,0,0):DP(1,1,0);
	printf("%lld\n",Ans);
	return 0;
}