h3mky0's blog

主に競プロについて書きます

ABC036 D - 塗り絵

問題リンク
atcoder.jp

木dpの練習

問題概要
N個の頂点を持つ木が与えられる. 各頂点を黒と白に塗ることができるが, 黒同士が隣接してはいけない. このとき色の塗り方が何通りあるか求める.

木dpは(treeの上を)トップダウンに数え上げる考え方(ととらえている)

dp[i][j] = 頂点iをjに塗った時の組み合わせ数(jが0なら白, jが1なら黒とする)

もし頂点iを白に塗るならば, その頂点の子は黒でも白でも問題ない.
もし頂点iを黒に塗るならば, その頂点の子はすべて白である必要がある.

これをもとに数え上げる

#include <bits/stdc++.h>
using namespace std;
#define MOD 1000000007
int N;
vector<vector<int>> G(100010);

ll dp[100010][2]; //0が白, 1が黒

void rec(int v, int prev = -1){

    ll sum_b = 1, sum_w = 1;
    for(int i=0; i<G[v].size(); i++){
        if(G[v][i] == prev) continue;

        rec(G[v][i], v);

        sum_b *= dp[G[v][i]][0]; 
        sum_w *= (dp[G[v][i]][1] + dp[G[v][i]][0]);
        sum_b %= MOD;
        sum_w %= MOD;
    }
    
    dp[v][0] = sum_w;
    dp[v][1] = sum_b;
}
int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    //cout << fixed << setprecision(15);
    
    cin >> N;
    REP(i,N-1){
        int a, b; cin >> a >> b;
        a--; b--;
        G[a].push_back(b);
        G[b].push_back(a);
    }
    REP(i,N){
        dp[i][0] = dp[i][1] = 1;
    }
    rec(0);
    cout << (dp[0][0]+dp[0][1])%MOD << endl; 
    return 0;
}