CF461B
设 \(f_{u,0/1}\) 表示以 \(u\) 为根的子树,\(u\) 所在的联通块内有 \(0/1\) 个黑点的方案数。
设 \(v\) 为 \(u\) 当前枚举到的儿子。
则转移方程为:
\(f_{u,1}=f_{u,1}\times(f_{v,0}+f_{v,1})+f_{u,0}\times f_{v,1}\)
\(f_{u,0}=f_{u,0}\times(f_{v,0}+f_{v,1})\)
初值为 \(f_{i,col_i}=1\),最终答案为 \(f_{1,1}\)。
时间复杂度 \(\mathcal O(n)\)。
Code:
#include <bits/stdc++.h>
using namespace std;
#define ll long long
const int N = 100005, mod = 1e9 + 7;
int n;
int col[N];
int head[N], ver[N*2], nxt[N*2], cnt;
ll f[N][2];
void add(int u, int v) {
ver[++cnt] = v, nxt[cnt] = head[u], head[u] = cnt;
}
void dfs(int u, int fa) {
f[u][col[u]] = 1;
for (int i = head[u], v; i; i = nxt[i]) {
v = ver[i];
if (v == fa) continue;
dfs(v, u);
f[u][1] = f[u][1] * f[v][1] % mod + f[u][1] * f[v][0] % mod + f[u][0] * f[v][1] % mod;
f[u][1] %= mod;
f[u][0] = f[u][0] * (f[v][1] + f[v][0]) % mod;
f[u][0] %= mod;
}
}
int main() {
scanf("%d", &n);
for (int i = 1, v; i < n; ++i) {
scanf("%d", &v);
add(i + 1, v + 1), add(v + 1, i + 1);
}
for (int i = 1; i <= n; ++i) scanf("%d", &col[i]);
dfs(1, 0);
printf("%lld", f[1][1]);
return 0;
}