[Python] 백준10830 행렬 제곱

16 Jun 2023 in Study on Algorithm

백준10830 행렬 제곱 링크

문제

예제 입력

---

코드

import sys
input = sys.stdin.readline

N, B = map(int, input().split())
arr = [list(map(int, input().split())) for _ in range(N)]

def matrix_mul(mat1, mat2):
    tmp = [[0]*N for _ in range(N)]
    for i in range(N):
        for j in range(N):
            for k in range(N):
                tmp[i][j] += mat1[i][k] * mat2[k][j] % 1000
    
    return tmp

def sol(b):
    if b == 1:
        return arr
    else :
        tmp = sol(b//2)
        if b % 2 == 0:
            return matrix_mul(tmp, tmp)
        else :
            return matrix_mul(matrix_mul(tmp, tmp), arr)
        
ans = sol(B)
for row in ans:
    for col in row:
        print(col % 1000, end=" ")
    print()

설명

파이썬을 통해서 사용자로부터 입력받아 분할 정복알고리즘을 사용하여 행렬 제곱를 구현했습니다.

---

정리

ans = arr.copy()
for _ in range(b-1):
    tmp = [[0]*n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            for k in range(n):
                tmp[i][j] += ans[i][k] * arr[k][j]
            tmp[i][j] = tmp[i][j] % 1000
    ans = tmp.copy()

for t in ans:
    print(*t)

처음에는 반복문으로 했지만 했던 연산을 계속해서 하기 때문에 시간 초과가 발생했습니다. 그렇기 때문에 분할 정복으로 나누고 기존 연산을 활용하여 시간을 줄였습니다.

---

결과