CFR 461 B. Appleman and Tree ( Tree DP )

Problem - 461B - Codeforces

題意：
給一顆樹，和每個節點是黑是白。求樹的分割組合數，滿足分割後的所有組合都恰有一個黒点，求其余数。

數據範圍：
節點數 2 ≤ N ≤ 1e5

解法：
樹形動態規劃寫起來真蛋疼。。。
將 0 作為根，轉換為有根樹上的問題。
dp[ i ][ j ] : 以 i 號節點為根的子樹，含 i 的分割無黒点 / 恰有一黒点，此時的分割組合數。

時間 / 空間複雜度：
O( N )

int N;
vi P;
vi X;

void init(){
    cin >> N;
    P = vi( N );
    for( int i = 1; i < N; ++i )
        cin >> P[ i ];
    X = vi( N );
    for( int i = 0; i < N; ++i )
        cin >> X[ i ];
}

vvi G;
vvi dp;

void dfs( int u, int fa ){
    if( ~dp[ u ][ 0 ] )
        return;
    dp[ u ][ not X[ u ] ] = 0;
    dp[ u ][ X[ u ] ] = 1;
    for( int v : G[ u ] )
        if( v != fa ){
            dfs( v, u );
            dp[ u ][ 1 ] = ( 1LL * dp[ u ][ 0 ] * dp[ v ][ 1 ] + 1LL * dp[ u ][ 1 ] * ( dp[ v ][ 0 ] + dp[ v ][ 1 ] ) ) % MOD7;
            dp[ u ][ 0 ] = 1LL * dp[ u ][ 0 ] * ( dp[ v ][ 0 ] + dp[ v ][ 1 ] ) % MOD7;
        } 
}

void preprocess(){
    G = vvi( N );
    for( int i = 1; i < N; ++i )
        G[ i ].emplace_back( P[ i ] ),
        G[ P[ i ] ].emplace_back( i );
    dp = vvi( N, vi( 2, -1 ) );
    dfs( 0, -1 );
}

void solve(){
    cout << dp[ 0 ][ 1 ] << endl;
}