[백준/Python] 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제곱한 결과를 출력한다.

풀이
풀이는 정확하게 잘 접근해서 결과도 잘 나오도록 출력했지만 아주 사소한 부분에서 시간초과가 계속 발생했다
범위가 크다보니 작은 연산속도 차이에도 시간초과가 발생할 수 있는 빡빡한 문제였다 ..

우선 두 행렬을 제곱하는 함수와 b를 분할해서 재귀를 호출하는 함수 두 개가 필요하다
행렬 A를 B 제곱하는 것은 $A^{B/2} \times A^{B/2}$와 동일하기 때문에 B가 1이 될 때 까지 재귀를 호출한다
만약 B가 홀수라면 $A^{B/2} \times A^{B/2} \times A$와 동일하므로 $A^{B/2}$끼리 행렬곱을 수행한 결과에 A를 한 번 더 행렬곱해줘야한다
이 과정에서 단순히 행렬을 두개로 분할하여 행렬곱을 수행하는 것 자체를 변수에 저장하고 호출을 한다면

ans = dotproduct(matmul(a,b//2), matmul(a,b//2))

안그래도 연산량이 많은 matmul 연산을 두 번 해야하기 때문에 시간 초과가 나기 쉽다
또한, 행렬 곱을 수행하는 함수에서 각 요소들을 계산할 때 1000을 모듈러연산을 해주어야 한다
요소를 계산하는 과정에서 곱하기 연산이 수행되기 때문에 값이 급격하게 커질 수 있으므로 계속해서 1000을 모듈러 연산을 수행해줌으로써 값의 크기를 줄일 수 있기 때문이라 추측하고 있었는데
챗지피티에 물어보니 그 추측이 맞았다
이런 사소한 부분때문에 코딩테스트에서 오래 헤매고 틀릴 수 있다고 생각하면 ,, 아찔하다

코드

import sys

n, b = map(int, sys.stdin.readline().split())
matrix = [list(map(int, sys.stdin.readline().split())) for _ in range(n)]

def dotproduct(a1, a2):
    result = [[0]*n for _ in range(n)]
    
    for i in range(n):
        for j in range(n):
            for k in range(n):
                result[i][j] += a1[i][k]*a2[k][j] % 1000
    return result

def matmul(a,b):
    if b == 1:
        return a
    ans = matmul(a,b//2)
    if b % 2 == 0:
        return dotproduct(ans, ans)
    else:
        return dotproduct(dotproduct(ans, ans), a)
    
answer = matmul(matrix,b)
for i in range(n):
    for j in range(n):
        answer[i][j] %= 1000

for elem in answer:
    print(*elem)