PS/Java

[백준] 10830. 행렬 제곱

siyamaki 2022. 3. 30. 17:27

크기가 N*N인 행렬 A가 주어진다. 이때, A의 B제곱을 구하는 프로그램을 작성하시오. 수가 매우 커질 수 있으니, A^B의 각 원소를 1,000으로 나눈 나머지를 출력한다.


입력

첫째 줄에 행렬의 크기 N과 B가 주어진다. (2 ≤ N ≤  5, 1 ≤ B ≤ 100,000,000,000)

둘째 줄부터 N개의 줄에 행렬의 각 원소가 주어진다. 행렬의 각 원소는 1,000보다 작거나 같은 자연수 또는 0이다.

출력

첫째 줄부터 N개의 줄에 걸쳐 행렬 A를 B제곱한 결과를 출력한다.


예제 입력 2 

3 3
1 2 3
4 5 6
7 8 9

예제 출력 2 

468 576 684
62 305 548
656 34 412

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {
    public static int N;
    public static long B;
    public static long[][] arr;
    public static void main(String[] args) throws Exception {

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        B = Long.parseLong(st.nextToken());

        arr = new long[N][N];

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

        long[][] res = matrixPow(arr, B); 

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

    public static long[][] matrixPow(long[][] tmp, Long exp) {
        if(exp == 1) {
            return tmp;
        }

        long[][] res = matrixPow(tmp, exp / 2);

        if(exp % 2 == 1) {
            return matMul(matMul(res, res), tmp);
        } else {
            return matMul(res, res);
        }
    }

    public static long[][] matMul(long[][] tmp, long[][] tmp1) {
        long[][] res = new long[N][N];

        for(int i = 0; i < N; i++) {
            for(int j = 0; j < N; j++) {
                long sum = 0;
                for(int k = 0; k < N; k++) {
                    sum += (tmp[i][k] * tmp1[k][j]) % 1000 ;
                }
                res[i][j] = sum;
            }
        }
        return res;
    }
}

matMul 함수로 행렬 간 제곱을 하는 함수를 만든다. 이 때 각 값을 더하기 전에 항상 1000으로 모듈러 연산을 해야한다.

public static long[][] matMul(long[][] tmp, long[][] tmp1) {
    long[][] res = new long[N][N];

    for(int i = 0; i < N; i++) {
        for(int j = 0; j < N; j++) {
            long sum = 0;
            for(int k = 0; k < N; k++) {
                sum += (tmp[i][k] * tmp1[k][j]) % 1000 ;
            }
            res[i][j] = sum;
        }
    }
    return res;
}

B제곱의 조건이 최대 10억번이므로 분할 정복 알고리즘을 사용하여야 한다. 큰 문제를 작은 문제로 쪼갠다음 계산하고 그 결과를 다시 합치는것이다.(https://ko.wikipedia.org/wiki/%EB%B6%84%ED%95%A0_%EC%A0%95%EB%B3%B5_%EC%95%8C%EA%B3%A0%EB%A6%AC%EC%A6%98)

 

public static long[][] matrixPow(long[][] tmp, Long exp) {
    if(exp == 1) {
        return tmp;
    }

    long[][] res = matrixPow(tmp, exp / 2);

    if(exp % 2 == 1) {
        return matMul(matMul(res, res), tmp);
    } else {
        return matMul(res, res);
    }
}

홀수 인 경우는 제곱을 한번 한 뒤 자기자신을 곱하면 된다.

짝수인 경우는 그냥 두 행렬의 곱셈을 수행하면 된다.