CF-461-B-树形DP

461-B 题目大意

给定一棵\(n\)个节点的树，节点编号从\(0\)开始，每个节点要么为白色要么为黑色，你需要删除一些边，使得剩下的各个连通块中有且仅有一个黑色节点。

问有多少种删边方案数，答案对\(10^9+1\)取模。

---

Solution

考虑树形DP，令\(dp[x][0/1]\)表示节点\(x\)属于无黑色节点/有黑色节点的连通块的方案数，节点\(y\)是\(x\)的儿子节点，初始化\(dp[x][col[x]]=1\)。

对于\(dp[x][1]\)，有三种情况：

• \(x\)属于有黑色节点的连通块，\(y\)属于有黑色节点的连通块，删去两点之间的边。
• \(x\)属于有黑色节点的连通块，\(y\)属于没有黑色节点的连通块，无需删边。
• \(x\)属于没有黑色节点的连通块，\(y\)属于有黑色节点的连通块，无需删边。

对于\(dp[x][0]\)，有两种情况：

• \(x\)属于没有黑色节点的连通块，\(y\)属于有黑色节点的连通块，删去两点之间的边。
• \(x\)属于没有黑色节点的连通块，\(y\)属于没有有黑色节点的连通块，无需删边。

那么转移方程如下：

\[dp[x][1]=dp[x][1]*dp[y][0]+dp[x][1]*dp[y][1]+dp[x][0]*dp[y][1] \]

\[dp[x][0]=dp[x][0]*dp[y][0]+dp[x][0]*dp[y][1] \]

要注意的是\(dp[x][1]\)的转移受\(dp[x][0]\)的影响，所以要先进行前者的转移，时间复杂度\(O(n)\)。

#include<bits/stdc++.h>
using namespace std;
using ll=long long;

const int mod=1e9+7;

void solve(){
    int n;
    cin>>n;
    vector<vector<int>> e(n);
    for(int i=1;i<n;i++){
        int x;
        cin>>x;
        e[x].push_back(i);
    }
    vector<array<ll,2>> dp(n);
    for(int i=0;i<n;i++){
        int col;
        cin>>col;
        dp[i][col]=1;
    }
    function<void(int)> dfs=[&](int x){
        for(auto y:e[x]){
            dfs(y);
            dp[x][1]=((dp[x][1]*dp[y][1]%mod)+(dp[x][1]*dp[y][0]%mod)+(dp[x][0]*dp[y][1]%mod))%mod;
            dp[x][0]=((dp[x][0]*dp[y][1]%mod)+(dp[x][0]*dp[y][0]%mod))%mod;
        }
    };
    dfs(0);
    cout<<dp[0][1];
}

int main(){
    ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
    //freopen("input.txt","r",stdin);
    //freopen("output.txt","w",stdout);
    int T=1;
    //cin>>T;
    while(T--){
        solve();
    }
    return 0;
}