[백준 10830] 행렬 제곱 Python

10830번 - 행렬 제곱

문제 이해

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

접근

행렬 곱을 정의한 후, 분할 정복을 이용한 거듭제곱을 사용한다.
단순 반복을 통해 $ O(N) $ 의 시간복잡도로 작성 할 시, 시간초과가 발생한다.
분할 정복을 이용한 거듭제곱

정답 코드

import sys
input = sys.stdin.readline


MOD = 1000
N, B = map(int, input().split())

M = [[*map(int, input().split())] for _ in range(N)]

# 행렬 곱
def matmul(matA, matB):
    ret = [[0]*N for _ in range(N)]
    for i in range(N):
        for j in range(N):
            for k in range(N):
                ret[i][j] += matA[i][k]*matB[k][j]
            ret[i][j] %= MOD
    return ret


result = [[0]*N for _ in range(N)]
for i in range(N):
    result[i][i] = 1
# fast power
while B > 0:
    if B % 2:
        result = matmul(result, M)
    M = matmul(M, M)
    B //= 2

for row in result:
    print(*row)