Codeforces 461B 树形 DP

题意

传送门 Codeforces 461B Appleman and Tree

题解

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

#include 
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;
}

你可能感兴趣的:(DP,算法)