[Baekjoon] 10830. 행렬 제곱

문제 설명

문제

크기가 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제곱한 결과를 출력한다.

예제 입력 1

2 5
1 2
3 4

예제 출력 1

69 558
337 406

예제 입력 2

3 3
1 2 3
4 5 6
7 8 9

예제 출력 2

468 576 684
62 305 548
656 34 412

예제 입력 3

5 10
1 0 0 0 1
1 0 0 0 1
1 0 0 0 1
1 0 0 0 1
1 0 0 0 1

예제 출력 3

512 0 0 0 512
512 0 0 0 512
512 0 0 0 512
512 0 0 0 512
512 0 0 0 512

문제 풀이

# 분할정복

분할 정복을 이용해서 행렬의 거듭제곱을 구하는 문제입니다.

이 문제를 풀기 위해서는 수의 거듭제곱을 분할정복을 이용해 푸는 방법과 행렬 곱을 구하는 법을 알아야 하는데요, 이 두 방법에 대한 문제와 풀이는 아래에서 확인하실 수 있습니다.

• [Baekjoon] 1629. 곱셈

• [Baekjoon] 2740. 행렬 곱셈

전체 코드입니다.

def cal_matrix_mul(A, B): # A*B
    C = [[0 for _ in range(len(B[0]))] for _ in range(len(A))]
    for i in range(len(A)):
        for j in range(len(B[0])):
            for k in range(len(A[0])):
                C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % 1000
    return C

def cal_matrix_power(A, b): # A^b
    if b in memo: X, Y = memo[b], I
    elif b % 2 == 0: X = Y = cal_matrix_power(A, b//2)
    else: X, Y = cal_matrix_power(A, b-1), A

    if b == 1 or b not in memo: memo[b] = cal_matrix_mul(X, Y)
    return memo[b]

n, b = map(int, input().split())
A = [list(map(int,input().split())) for _ in range(n)]
I = [[0]*k + [1] + [0]*(len(A)-k-1) for k in range(len(A))] # 항등 행렬
memo = {1: A}
res = cal_matrix_power(A, b)
print(*[" ".join(map(str,r)) for r in res],sep='\n')

line 10에서 바로 memo[b]를 반환하지 않고 항등 행렬 I와 곱한 이유는 b = 1 인 경우에 각 원소를 1000으로 나누기 위함입니다.