[Algorithm] 백준 알고리즘 9521번 색칠 공부

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

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.StringTokenizer;

public class Baekjoon9521 {
//public class Main {
	
	static long MOD = 1000000007;
	static long[] dp;
	
	public static void main(String[] args) 
			throws IOException
	{
		
		InputStreamReader fr = new InputStreamReader(System.in);
//		FileReader fr =new FileReader("D:\\dev\\worksspace\\Solution\\src\\Baekjoon9521");
		BufferedReader br = new BufferedReader(fr);
		BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
		String s;
		StringTokenizer st;
//		int totTestNum = Integer.parseInt(br.readLine());
//		int curTestNum = 0;
		
		while ((s=br.readLine())!=null) {
//			curTestNum++;
//			if(curTestNum > totTestNum) break;
			
			st = new StringTokenizer(s);
			int N = Integer.parseInt(st.nextToken());
			int K = Integer.parseInt(st.nextToken());
			dp = new long[N+1];
			dp[0] = 1;
			for (int i = 1; i <= N; i++) {
				dp[i] = dp[i-1] * (K-1) % MOD;
			}
			
			st = new StringTokenizer(br.readLine());
			int[] arr = new int[N+1];
			for (int i = 1; i <= N; i++) {
				arr[i] = Integer.parseInt(st.nextToken());
			}
			
			//알고리즘 시작
			int[] visitCnt = new int[N+1];
			int[] visitGrp = new int[N+1];
			int grp = 0;
			long answer = 1;
			
			for (int i = 1; i <= N; i++) {
				if(visitGrp[i] < 1){
					grp++;
					visitGrp[i] = grp;
					visitCnt[i] = 1;
					
					int curVisitCnt = 1;
					int nextNode = i;
					while (true) {
						nextNode = arr[nextNode];
						
						if(visitGrp[nextNode] < 1){
							visitGrp[nextNode] = grp;
							visitCnt[nextNode] = ++curVisitCnt; 
							
						}else{
							if(visitGrp[nextNode] == grp){
								//1. 같은 그룹일 경우
								answer = (answer * K) %MOD;
								answer = (answer * cycleAnswer(curVisitCnt-visitCnt[nextNode])) %MOD;
								answer = (answer * dp[visitCnt[nextNode]-1]) %MOD;
							}else{
								//2. 다른 그룹일 경우
								answer = (answer * dp[curVisitCnt]) %MOD;
							}
							break;
						}
					}
				}
			}
//			bw.write(answer + "\n");
			System.out.println(answer);
		}
		
		bw.flush();
		bw.close();
		br.close();
	}

	private static long cycleAnswer(int cnt) {
		
		if(cnt < 1) return 1;
		long pm = 1;
		long result = 0;
		
		for (int i = cnt; i >0; i--) {
			result += pm*dp[i];
			if(result < 0) result += MOD;
			result = result % MOD;
			pm *= -1;
		}
		
		return result;
	}
}