CF - 461B

문제링크

링크텍스트

제한 조건

time limit per test2 seconds, memory limit per test 256 megabytes

문제 스케치

tree DP인거 같은데 세는 방식이 어려웠다. 일단 일차원 DP로 접근해 봤는데 검은색 node를 포함하는지 안 하는지 경계를 찾기가 매우 힘들었다. 그래서 이차원 DP 첫번째는 node 두번째는 0,1 해당 노드에 대해서 검은색 노드가 포함되었는지를 확인하게 해주었다. dp initialize는 일단 기본적으로 해당 노드가 검은색인지 하얀색인지 leaf node까지 내려가면서 initialize 해주었고 이제 경우의 수를 세는게 문제였는데 진짜 어려워서 머리가 터질뻔 했다 분명 DP는 맞는데 이걸 어떻게 해야하는지 몰라가지고 고생했다
일단 검은색은 혼자 독립적으로 존재 가능하다. 간선을 자를 때 검은색 노드가 자식 쪽에 있을 때 간선을 안 자를 때 부모 노드에 대해서 검은 노드 있고 간선 자를 때와 안 자를 때 자식에 검은 노드 없는 경우 + 부모 쪽에 검은노드가 없을 때 자식 쪽에 검은 노드 있는 경우를 해서 dp를 업데이트해 나가는건데 진짜 너무 어렵다...

#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
#include <queue>
#include <stack>
#include <deque>
#include <map>
#include <unordered_map>
#include <set>
#include <unordered_set>
#include <cmath>
#include <numeric>
#include <limits>
#include <sstream>
#include <iomanip> 

#define INF 0x3f3f3f3f // 경우에 따라 다르게
#define MOD 1000000007
// long long 일 1e18

using namespace std;

int n;
vector<int> edge[100001];
int color[100001];
long long dp[100001][2];


void dfs(int v, int parent)
{
    //cout <<  v <<" "<<parent<<"\n";
    if(color[v] == 1)
        dp[v][1] +=1; 
    else
        dp[v][0] +=1;
    
    for(auto c : edge[v])
    {
        if(c==parent)
            continue;
        dfs(c,v);


        // 일차원 dp 가 아니긴 한데 
        // dp[v]
        long long w = dp[v][0];
        long long b = dp[v][1];
    
        // 간선을 자를 때 
        // 자식 쪽에 검은노드 한 개 일때 
        
    
        // 간선을 안 자를 때 
        // v쪽에 검은 노드 있고 없고 , c쪽에 검은 노드 있고 없고
  
        dp[v][0] =  (w * (dp[c][1] + dp[c][0] ) )% MOD; // 간선 자를 때 + 안 자를 때 c에 검은 없는경우
        dp[v][1] = (b * (dp[c][1] + dp[c][0])  + w * dp[c][1] ) % MOD;

    }



}

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

    cin >> n ;
    for(int i =  1; i<n; i++)
    {
        int p ;
        cin >> p;
        edge[p].push_back(i);
        edge[i].push_back(p);
    }


    for(int i =0 ;i<n; i++)
    {
        int c ;
        cin >> c; 
        color[i] = c; 
    }
    

    dfs(0,-1);

    
    // for(int i = 0; i< 2; i++)
    // {
    //     for(int j =0 ; j< n; j++)
    //     {
    //         cout << dp[j][i] << " ";
    //     }
    //     cout <<"\n";
    // }

    cout << dp[0][1] % MOD<<"\n";

}