【树形DP】CODEFORCES 461B Appleman and Tree

通道

题意：一棵树上有K个黑色节点，剩余节点都为白色，将其划分成K个子树，使得每棵树上都只有1个黑色节点，共有多少种划分方案

思路：dp[i][0]和dp[i][1]分别表示i子树所在联通块不存在黑节点和已经存在一个黑节点的方案数

代码：

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<vector>
#include<math.h>
#include<queue>
using namespace std;
const int mod=1e9+7;
int n,color[100005];
vector <int> map[100005];
long long dp[100005][2];
void dfs(int now)
{
    dp[now][0]=1;
    dp[now][1]=0;
    for(int i=0;i<map[now].size();i++)
    {
        int to=map[now][i];
        dfs(to);
        dp[now][1]=(dp[now][1]*dp[to][0]+dp[now][0]*dp[to][1])%mod;
        dp[now][0]=dp[now][0]*dp[to][0]%mod;
    }
    if(color[now]==1)
    dp[now][1]=dp[now][0];
    else
    dp[now][0]=(dp[now][0]+dp[now][1])%mod;
}
int main()
{
    int n;
    cin>>n;
    for(int i=1;i<n;i++)
    {
        int a;
        scanf("%d",&a);
        map[a].push_back(i);
    }
    for(int i=0;i<n;i++)
    scanf("%d",&color[i]);
    dfs(0);
    long long ans=dp[0][1];
    cout<<ans<<endl;
    return 0;
}