백준 #10830 행렬 제곱 (파이썬)

오늘의 한마디
좁은 부분을 검사한 결과 완전히 맞다는 확신이 든다면, 좀 더 넓게. 전체적인 코드 진행을 보아라.

문제 링크

선형대수학을 공부하고 있는 기념으로, 선형대수학 기본 문제를 풀어줬다.
그리고 솔브닥의 CLASS 4에도 수록되어 있다.

• mmul = matrix_multiply
• mpow = matrix_power

이 두 함수도 클리셰라 외울만 하다.

80%에서 '틀렸습니다'가 나오길래 계속 고민해도 알고리즘이 너무 정확해 오류를 못 찾았는데, 문제가 엉뚱한 엣지 케이스에서 있었다.

import sys

N, B = map(int, sys.stdin.readline().rstrip().split())
arr = []
for _ in range(N):
    arr.append(list(map(int, sys.stdin.readline().rstrip().split())))

# Edge case
for i in range(N):
    for j in range(N):
        if arr[i][j] == 1000:
            arr[i][j] = 0
    
def mmul(A, B):
    C = [[0] * len(B[0]) for _ in range(len(A))]
    for i in range(len(C)):
        for j in range(len(C[0])):
            for k in range(len(A[0])):
                C[i][j] += A[i][k] * B[k][j]
            C[i][j] %= 1000
    return C

def mpow(A, n):
    if n == 1:
        return A
    A = mpow(A, n//2)
    A = mmul(A, A)
    if n%2 == 1:
        return mmul(A, arr)
    return A

arr = mpow(arr, B)
for i in range(N):
    print(' '.join(map(str, arr[i])))