Codeforces 461B 树形 DP

题意

传送门 Codeforces 461B Appleman and Tree

题解

$d{p}_{v,k}$ 代表以节点 $v$ 为根的子树中，包含了 $v$ 的联通分量是否存在一个黑色节点 ，同时其余联通分量仅包含一个黑色节点情况下，划分方案的数量。DFS 求解，对于每一条连向子节点的边，考虑是否删去这条边，进行状态转移即可。总时间复杂度 $O(n)$。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
constexpr int MOD = 1e9 + 7;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    cin >> n;
    vector<vector<int>> g(n);
    for (int i = 0; i < n - 1; ++i) {
        int p;
        cin >> p;
        g[p].push_back(i + 1);
    }
    vector<int> x(n);
    for (int i = 0; i < n; ++i) {
        cin >> x[i];
    }
    vector<vector<int>> dp(n, vector<int>(2));
    function<void(int)> dfs = [&](int v) {
        dp[v][x[v]] = 1;
        for (int u : g[v]) {
            dfs(u);
            auto tmp = dp[v];
            dp[v][0] = dp[v][1] = 0;
            for (int k = 0; k < 2; ++k) {
                (dp[v][k] += (ll)tmp[k] * dp[u][1] % MOD) %= MOD;
            }
            for (int k = 0; k < 2; ++k) {
                (dp[v][k] += (ll)tmp[k] * dp[u][0] % MOD) %= MOD;
            }
            (dp[v][1] += (ll)tmp[0] * dp[u][1] % MOD) %= MOD;
        }
    };
    dfs(0);
    cout << dp[0][1] << '\n';

    return 0;
}