[BOJ] 백준 17401 일하는 세포 (자바)

자바 17401 일하는 세포 (링크)

풀이

행렬의 거듭제곱을 이용해 한 사이클에 해당하는 행렬을 구하고 D초 후의 결과를 계산하는 문제이다. 문제 자체는 어려워 보이지만 단순한 행렬 곱셈 문제와 다를 게 없다.

N*N 배열을 만들고 각 초마다 바뀌는 통로를 저장한 뒤 해당 배열을 계속 곱해 주자. 처음에 단위 행렬 배열을 만들고 곱해 나가면 된다.
그리고 주기가 T초라면, D/T 만큼 한 사이클에 해당하는 행렬을 곱해 주면 된다.
예를 들어 사이클이 3초이고, 총 10초동안 이동한다면 3번 사이클행렬을 곱하고 남은 1번을 마저 곱해 주면 되는 것이다.

하지만 문제를 보면, D의 범위가 10^9 까지이기 때문에 한 T가 1초일 경우 행렬을 10^9번 곱해야 한다. 시간초과를 방지하기 위해 다음과 같이 bin배열에 사이클 행렬을 2^n 제곱한 것을 저장한 뒤 꺼내어 곱할 것이다. (2의 30승이 D의 범위보다 크기에 길이를 31로 한다)

최종 코드 및 결과

import java.io.*;
import java.util.*;

public class Main {
    static int T, N, D, INF = 1000000007;
    public static void main(String[] args) throws IOException {
    	BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    	StringTokenizer st = new StringTokenizer(br.readLine());
    	T = Integer.parseInt(st.nextToken()); // 주기
    	N = Integer.parseInt(st.nextToken()); // 거점수
    	D = Integer.parseInt(st.nextToken()); // 시간
    	int pow = D/T; // 사이클을 반복할 횟수
		int[][] cycle = new int[N][N];	// 한 사이클에 해당하는 행렬
		int[][] res = new int[N][N];	// 답을 담을 행렬
		int[][] remain = new int[N][N];	// 주기를 전부 돌고 남은 이동을 담을 행렬
		int[][][] bin = new int[31][N][N];	// 2^n만큼 사이클을 거듭제곱한 행렬을 담음
    	
		// cycle, res, remain 단위행렬로 초기화
    	for(int i = 0; i < N; i++) remain[i][i] = res[i][i] = cycle[i][i] = 1;
    	for(int t = 0; t < T; t++) {
    		int M = Integer.parseInt(br.readLine());
    		int[][] matrix = new int[N][N]; // t초에 해당하는 통로를 담을 행렬
    		for(int i = 0; i < M; i++) {
    			st = new StringTokenizer(br.readLine());
    			int a = Integer.parseInt(st.nextToken());
    			int b = Integer.parseInt(st.nextToken());
    			int c = Integer.parseInt(st.nextToken());
    			matrix[a-1][b-1] = c; // matrix[시작][끝] = 통로개수
    		}
    		if (t < D%T) remain = mult(remain, matrix); // t < D%T 이면 remain에 곱하기
    		cycle = mult(cycle, matrix); // 한 사이클 행렬을 구하기 위해 계속 cycle에 곱하기
    	}
    	bin[0] = cycle; // 2^0 = 1 이므로 사이클을 1번 거듭제곱한 것을 넣어줌
    	
    	// idx 1에는 cycle^(2^1) 저장, idx 2에는 cycle^(2^2) 저장 ... 
    	for(int i = 1; i < 31; i++) bin[i] = (cycle = mult(cycle, cycle)).clone();
    	
    	// 30부터 0까지 pow >= 2^i 이면 pow - 2^i 후 res에 bin[i] 곱하기
    	for(int i = 30; i >= 0 && pow != 0; i--) {
    		if (pow < (1<<i)) continue;
    		pow -= 1<<i;
    		res = mult(res, bin[i]);
    	}
    	res = mult(res, remain); // 남은 이동 곱하기
    	StringBuilder sb = new StringBuilder();
    	for(int i = 0; i < N; i++) {
    		for(int j = 0; j < N; j++) sb.append(res[i][j]).append(" ");
    		sb.append("\n");
    	}
    	System.out.print(sb);
    }
    static int[][] mult(int[][] m1, int[][] m2) { // 행렬 곱셈 함수
        int[][] res = new int[N][N];
        for(int i = 0; i < N; i++) for(int j = 0; j < N; j++) for(int k = 0; k < N; k++) {
    		res[i][j] += ((long)m1[i][k]*m2[k][j])%INF;
    		res[i][j] %= INF;
        }
        return res;
    }
}