Appleman and Tree (Codeforces 461B)
Tags: DP, trees

Appleman has a tree with n vertices. Some of the vertices (at least one) are colored black and other vertices are colored white.Consider a set consisting of k (0 ≤ k < n) edges of Appleman's tree. If Appleman deletes these edges from the tree, then it will split into (k + 1) parts. Note, that each part will be a tree with colored vertices.Now Appleman wonders, what is the number of sets splitting the tree in such a way that each resulting part will have exactly one black vertex? Find this number modulo 1000000007 (109 + 7).

Input: The first line contains an integer n (2  ≤ n ≤ 105) — the number of tree vertices. The second line contains the description of the tree: n - 1 integers p0, p1, ..., pn - 2 (0 ≤ pi ≤ i). Where pi means that there is an edge connecting vertex (i + 1) of the tree and vertex pi. Consider tree vertices are numbered from 0 to n - 1.The third line contains the description of the colors of the vertices: n integers x0, x1, ..., xn - 1 (xi is either 0 or 1). If xi is equal to 1, vertex i is colored black. Otherwise, vertex i is colored white.

Output: Output a single integer — the number of ways to split the tree modulo 1000000007 (109 + 7).


Sample Inputs:
3
0 0
0 1 1

6
0 1 1 0 4
1 1 0 0 1 0

10
0 1 2 1 4 4 4 0 8
0 0 0 1 0 1 1 0 0 1

Sample Outputs:
2

1

27

Solution (Python):
MOD = 1000000007
n = int(input())
p = [int(x) for x in input().split()]
x = [int(x) for x in input().split()]
children = [[] for x in range(n)]
for i in range(1, n):
    children[p[i - 1]].append(i)
count = [(0, 0) for i in range(n)]
for i in reversed(range(n)):
    prod = 1
    for ch in children[i]:
        prod *= count[ch][0] + count[ch][1]
    if x[i]:
        count[i] = (0, prod % MOD)
    else:
        tot = 0
        for ch in children[i]:
            cur = count[ch][1] * prod // (count[ch][0] + count[ch][1])
            tot += cur
        count[i] = (prod % MOD, tot % MOD)
print(count[0][1])

Solution (Python):
import threading
import sys
sys.setrecursionlimit(3 * 10 ** 5 + 200)
threading.stack_size(3 * 10 ** 5 + 100)
MOD = 10 ** 9 + 7
(dp0u, dp1u) = (-1, -1)

def dfs(v):
    if c[v]:
        dp0v = 0
        dp1v = 1
        for u in edges[v]:
            (dp0u, dp1u) = (yield u)
            dp1v *= dp0u + dp1u
            dp1v %= MOD
    else:
        dp0v = 1
        dp1v = 0
        for u in edges[v]:
            (dp0u, dp1u) = (yield u)
            dp1v *= dp0u + dp1u
            dp1v += dp0v * dp1u
            dp0v *= dp0u + dp1u
            dp0v %= MOD
            dp1v %= MOD
    yield v
    yield (dp0v, dp1v)
n = int(input())
(*p,) = [int(x) for x in input().split()]
(*c,) = [int(x) for x in input().split()]
edges = [[] for _ in range(n)]
for i in range(n - 1):
    edges[p[i]].append(i + 1)
dp = [[0] * n, [0] * n]
for v in range(n)[::-1]:
    gen = dfs(v)
    u = gen.send(None)
    while u != v:
        u = gen.send((dp[0][u], dp[1][u]))
    (dp[0][v], dp[1][v]) = gen.send(None)
print(dp[1][0])

Solution (Python):
MOD = 1000000007
n = int(input())
p = [int(x) for x in input().split()]
x = [int(x) for x in input().split()]
children = [[] for x in range(n)]
for i in range(1, n):
    children[p[i - 1]].append(i)
count = [(0, 0) for i in range(n)]
for i in reversed(range(n)):
    prod = 1
    for ch in children[i]:
        prod *= count[ch][0] + count[ch][1]
    if x[i]:
        count[i] = (0, prod % MOD)
    else:
        tot = 0
        for ch in children[i]:
            cur = count[ch][1] * prod // (count[ch][0] + count[ch][1])
            tot += cur
        count[i] = (prod % MOD, tot % MOD)
print(count[0][1])