461B

Description

A tree has $n$($1 \leq n \leq 10^5 $) vertices, the color of each of which is either black or white. Now, we need to split the tree into many parts, each of which has exactly one black vertex. Output how many ways to split the tree.

Tutorial

This is a dynamic programming problem on a tree. Define dp[i][j][0] as the number of ways to split the subtree whose root is $i$, and only have its first $j$ subtrees (the subtrees whose roots are $i$'s direct children), just take the other children as not exist, and all the vertices in the part which contains vertex $i$ is white. (all the other parts in this subtree has exactly one black vertex.) Define dp[i][j][1] is the same. The only difference is that there is exactly one black vertex in the part contains vertex $i$ is white. (all the parts in this subtree has exactly one black vertex.)

dp[i][0][0]=1 and dp[i][0][1]=0 if the root is white, otherwise their values are opposite.

dp[i][j][0] = dp[i][j - 1][0] * dp[v][n_v][0];

dp[i][j][1] = dp[i][j - 1][0] * dp[v][n_v][1] + dp[i][j - 1][1] * dp[v][n_v][0];

The two equation above means do not split $i$ and $v$. ($v$ is $i$'s $j$th child. $n_v$ is the number of children $v$ has.)

dp[i][j][0] += dp[i][j - 1][0] * dp[v][n_v][1];

dp[i][j][1] += dp[i][j - 1][1] * dp[v][n_v][1];

The two addition above means we split $i$ and $v$. When calculating dp[x][j][y], we need dp[x][j - 1][y]. We can delete the second dimension, just replace the old value with the new as $j$ increases is ok.

There are many things I learned from this problem.

1. About the defining the max value of n.

     const int MAX_N = int(1e5) + 10;
   

2. Iterating all the scenarios and use conditions to control what to do. The following state transition controll is a good example.

     for (int i = 0; i < 2; i++)
     {
         for (int j = 0; j < 2; j++)
         {
             if (i + j < 2)
                 cur[i + j] = (cur[i + j] + 1LL * last[i] * dp[v][j] % MOD) % MOD;
             if (j == 1)
                 cur[i] = (cur[i] + 1LL * last[i] * dp[v][j] % MOD) % MOD;
         }
     }
   

3. A way to controll DFS.

     DFS()
     {
         for (all the children)
         {
             DFS();
         }
         do the work();
     }
   

4. The connection of dynamic programming of two different data structures of storing trees. One is as a graph. The other is "left son right brother". The thought of going through the children one by one and see the past children as a whole is very similar with the thought of DP on "left son right brother".

Solution

#include <cstdio>
#include <vector>
using namespace std;

const int MAX_N = int(1e5) + 10;
const int MOD = int(1e9) + 7;

int n;
vector<int> edge[MAX_N];
int color[MAX_N];
int dp[MAX_N][2]; //0 white, 1 black

void input()
{
    scanf("%d", &n);
    for (int i = 0; i < n - 1; i++)
    {
        int p;
        scanf("%d", &p);
        edge[i + 1].push_back(p);
        edge[p].push_back(i + 1);
    }
    for (int i = 0; i < n; i++)
    {
        scanf("%d", &color[i]);
    }
}

void dfs(int u, int parent)
{
    for (vector<int>::iterator it = edge[u].begin(); it != edge[u].end(); it++)
    {
        int v = *it;
        if (v == parent)
            continue;
        dfs(v, u);
    }

    vector<int> last(2, 0);
    last[color[u]] = 1;
    for (vector<int>::iterator it = edge[u].begin(); it != edge[u].end(); it++)
    {
        int v = *it;
        if (v == parent)
            continue;
        vector<int> cur(2, 0);
        for (int i = 0; i < 2; i++)
        {
            for (int j = 0; j < 2; j++)
            {
                if (i + j < 2)
                    cur[i + j] = (cur[i + j] + 1LL * last[i] * dp[v][j] % MOD) % MOD;
                if (j == 1)
                    cur[i] = (cur[i] + 1LL * last[i] * dp[v][j] % MOD) % MOD;
            }
        }
        last = cur;
    }

    for (int i = 0; i < 2; i++)
    {
        dp[u][i] = last[i];
    }
}

int main()
{
    input();
    dfs(0, -1);
    printf("%d\n", dp[0][1]);
    return 0;
}