CF 461B Appleman and Tree

https://vjudge.net/problem/CodeForces-461B

题目

某人有一颗树，上面有些节点是黑色的，他可以去掉一些边，让这颗树变成森林，并且每一个树都有且只有1个黑色节点，问有多少种选择边的方法。

$nleqslant 10^5$

题解

首先把树转换为有根树

然后dp[x][1/0]表示节点x上面那条边下面的子树只有一个黑节点或没有黑节点的方案数

规定边可以断开：

假设x没有颜色，v0表示从x开始的子树都是白色的方案个数，v1表示x的子树有一个黑色节点的方案个数

那么[egin{array}{cc}v1=&dp[s_1][1]*dp[s_2][0]*...*dp[s_n][0]\+&dp[s_1][0]*dp[s_2][1]*dp[s_3][0]*...*dp[s_n][0]\+&...\+&dp[s_1][0]*dp[s_2][0]*...*dp[s_{n-1}][0]*dp[s_n][1]end{array}]

$v0=dp[s_1][0]*dp[s_2][0]*...*dp[s_n][0]$

如果x是白色，$dp[x][1]=v1$，$dp[x][0]=v1+v0$（因为边可以断开）

如果x是黑色，$dp[x][1]=v0$，$dp[x][0]=v0$（因为即使断开子树也只能有一个黑节点）

但是为了计算v1需要花$mathcal{O}(n^2)$的时间，肯定超时

[egin{array}{cccccc}&1& imes&0& imes&0\+&0& imes&1& imes&0\+&0& imes&0& imes&1end{array}]

其实可以递推计算v1：我们可以从左上角的$k imes k$的正方形递推到$(k+1) imes(k+1)$的正方形

那么计算v1只需要花$mathcal{O}(n)$，每个点最多访问两次，时间复杂度$mathcal{O}(n)$

AC代码

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<vector>
#include<map>
#define REP(i,a,b) for(register int i=(a); i<(b); i++)
#define REPE(i,a,b) for(register int i=(a); i<=(b); i++)
#define PERE(i,a,b) for(register int i=(a); i>=(b); i--)
using namespace std;
typedef long long ll;
#define MAXN 100007
int n;
int hd[MAXN], nxt[MAXN*2], to[MAXN*2], m;
bool b[MAXN];
inline void adde(int a, int b) {
	nxt[m]=hd[a]; hd[a]=m; to[m]=b; m++;
}
#define MO (int(1e9)+7)
ll dp[MAXN][2];
bool vis[MAXN];
void dfs(int x) {
	ll v1=0, v0=1;
	vis[x]=1;
	int cnt=0;
	for(int i=hd[x]; ~i; i=nxt[i]) if(!vis[to[i]]) {
		cnt++;
		int o=to[i];
		dfs(o);
		v1=v1*dp[o][0]%MO;
		v1=(v1+v0*dp[o][1])%MO;
		v0=v0*dp[o][0]%MO;
	}
	if(b[x]) {
		dp[x][1]=v0;
		dp[x][0]=v0;
	} else {
		dp[x][1]=v1;
		dp[x][0]=(v1+v0)%MO;
	}
}
int main() {
	scanf("%d", &n); m=0;
	memset(hd,-1,sizeof(int)*n);
	memset(vis,0,sizeof(bool)*n);
	REP(i,1,n) {
		int x;
		scanf("%d", &x);
		adde(i,x);
		adde(x,i);
	}
	REP(i,0,n) {
		int t; scanf("%d", &t);
		b[i]=t;
	}
	dfs(0);
	printf("%lld
", dp[0][1]);
}