백준 10830번 : 행렬 제곱 with Java

import java.io.*;
import java.util.Stack;
import java.util.StringTokenizer;

public class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
    static StringTokenizer st = null;

    public static void main(String[] args) throws IOException {
        int n;
        long b;
        long[][] matrix;
        long[][] result;
        Stack<Integer> stack = new Stack<>();

        st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());
        b = Long.parseLong(st.nextToken());
        matrix = new long[n][n];
        result = new long[n][n];

        for (int i = 0; i < n; i++) {
            st = new StringTokenizer(br.readLine());
            for (int j = 0; j < n; j++) {
                matrix[i][j] = Long.parseLong(st.nextToken());
                result[i][j] = matrix[i][j];
            }
        }

        while (b > 0)
        {
            if (b % 2 == 1)
                stack.add(1);
            else
                stack.add(0);
            b /= 2;
        }

        stack.pop();

        while (!stack.isEmpty()) {
            long bit = stack.pop();
            result = mul(result, result, n);
            if (bit != 0) {
                result = mul(result, matrix, n);
            }
        }

        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                System.out.print(result[i][j] % 1000 + " ");
            }
            System.out.println();
        }
    }

    private static long[][] mul(long[][] matrix1, long[][] matrix2, int n) {
        long[][] n_matrix;

        n_matrix = new long[n][n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                long result = 0;
                for (int k = 0; k < n; k++) {
                    result += matrix1[i][k] * matrix2[k][j];
                    result %= 1000;
                }
                n_matrix[i][j] = result;
            }
        }
        return (n_matrix);
    }
}

문제에서 확인할 수 있듯이, 최대 행렬의 1000억 제곱까지 구해야 한다.

 

따라서 단순 Brute Force로 풀게 되면 시간이 팡팡 터지는 문제다.

그럼 어떻게 시간을 줄일 수 있을까?

개념은 정수의 N 거듭제곱 빠르게 구하기와 동일하다.

 

위 표는 11승을 구하는 방법이다.

 

11을 2진수로 표현하면 1011 이고, 이를 행렬 곱에 적용한다면 1일때 제곱을 한 뒤에 입력받은 행렬을 곱해준다.

0일때는 제곱만 한다.

 

조건은 매우 간단하다.

 

입력받은 B를 나눠서 2진수를 Stack에 담아주게 되면 최상위 비트부터 뽑아서 확인할 수 있다.

 

        while (b > 0)
        {
            if (b % 2 == 1)
                stack.add(1);
            else
                stack.add(0);
            b /= 2;
        }

다음 스택에서 무조건 한 개 꺼내준다.

 

그 이유는 다음과 같다.  =>   문제에서 주어지는 값이 최상위 비트가 항상 1임을 보장한다.

따라서 최상위 비트 계산을 할 필요없이, 주어지는 입력 행렬로 시작하면 된다.