백준 9521번 - 색칠 공부

https://www.acmicpc.net/problem/9521

처음 이 문제를 접하고, 생각보다 쉽게 해결될 수 있다고 생각했는데 왜 정답률이 22프로인지를 알 수 있던 문제다.

사이클 개념과 DP문제가 혼합되어 있어서 더욱 어렵게 느껴졌다.

먼저 문제를 해결하는 방식은 다음과 같다.

1) 사이클에 N개의 정점이 존재 할 때 N개의 정점을 색칠할 수 있는 경우의 수를 구한다.

2) 사이클에 포함되지 않는 정점들을 모두 K-1씩 곱해준다.

기존 사이클의 개수를 구하는 방식은 DFS 문제의 9486번(텀프로젝트)에 개시했다.

여기서 중요한 점은 사이클이 N개가 존재할 때

1번째 정점과 N-2번째 정점이 같은색인 경우와 다른 색인 경우 2가지가 존재한다.

예를 들어 1->2->3->4->1 과 같은 사이클이 있을 때

4개의 정점을 색칠하는 경우의 수를 구하려면 1번과 3번이 같은 색일 때와 다른 색일 때를 더해줘야 한다는 것이다.

즉, 점화식은 다음과 같다.

color[i] = color[i-2] * (K-1) + color[i-1] * (K-2)

(사이클에서 i개의 정점이 있을 때, 1번째 정점과 i-2번째 정점 색 같을 때 + 1번째 정점과 i-2번째 정점 색 다를 때)

사이클의 개수에 따라서 경우의 수를 모두 곱해주면 되며

전체 정점에서 사이클의 개수를 뺀 모든 정점은 다른 정점과 연결이 되어있기 때문에

연결된 정점과 다른 K-1가지 색을 계속해서 곱해주면 된다.

#include <iostream>
using namespace std;
#define inf 1000000007;
 
long ans=1;
int N;
long K;
 
int start_point=0; //사이클 방문 시 시작 정점
long color[1000001]={0}; //사이클 개수 별 경우의 수
int first_visit[1000001]={0}; //사이클 방문 시 노드 별 시작 정점
int visited[1000001]={0}; //사이클 방문 시 방문 순번
int f[1000001]={0};
 
int dfs(int current,int cnt){
    
    if(visited[current]){
        
        //기존 사이클과 다른 정점 방문 시 (사이클 x)
        if(first_visit[current]!=start_point)
            return 0;
        
        return cnt-visited[current];
    }
    
    visited[current]=cnt;
    first_visit[current]=start_point;
    return dfs(f[current],cnt+1);
}
 
int main(){
 
    cin>>N>>K;
 
    for(int i=1;i<=N;i++){
        scanf("%d",&f[i]);
    }
    
    color[0]=1;
    color[1]=K;
    color[2]=(K*(K-1))%inf;
    color[3]=(K*(K-1)*(K-2))%inf;
    
    for(int i=4;i<=N;i++){
        
        //1번째와 i-2가 같은색일 때 + 1번째와 i-2가 다른색일 때
        color[i]=color[i-2]*(K-1)+color[i-1]*(K-2);
        color[i]%=inf;
    }
    
    int single_cnt=N; //사이클이 아닌 정점 개수
    int cycle_cnt=0; //사이클에 포함되는 정점 개수
 
    for(int i=1;i<=N;i++){
        if(!visited[i]){
 
            start_point=i;
            cycle_cnt=dfs(i,1);
 
            ans*=color[cycle_cnt];
            ans%=inf;
            single_cnt-=cycle_cnt;
        }
    }
    
    for(int i=1;i<=single_cnt;i++){
        ans*=(K-1);
        ans%=inf;
    }
 
    cout<<ans<<endl;
    return 0;
}