BOJ 10830: 행렬 제곱 - 파이썬(Python)

https://www.acmicpc.net/problem/10830

<접근 방법>

크기가 $N \times N$인 행렬 $A$의 $B$제곱인 $A^B$를 구하는 문제이다.

 

단순 곱을 하게 된다면 행렬 곱의 시간복잡도 $O(N^3)$을 B번 진행하게 되므로 $O(n^3 \times B)$가 될 것이다. 이는 제한 시간인 1초 이내로 실행할 수 없는 시간복잡도이다.

 

그렇다면 어떻게해야 시간복잡도를 줄일 수 있을까? "BOJ 1629: 곱셈" 문제에서 사용했던 "분할 정복"을 적용한 곱셈을 행렬 버전으로 적용하면 될 것 같다고 생각할 수 있을 것이다. 

 

즉, 아래의 식을 이용하면 답을 구할 수 있다.

$A^B = A^\frac{B}{2} \times A^\frac{B}{2}$
$A^B = A \times A^{floor(\frac{B}{2})} \times A^{floor(\frac{B}{2})}$

 

그렇다면 시간복잡도는 어떻게 될까? 위의 식을 살펴보면 지수 B의 값이 절반으로 줄어든다는 것을 알 수 있다. 즉, 아래의 재귀식이 적용될 것이다. $T(B)$는 $B$제곱을 구하는데 걸리는 총 시간, $O(N^3)$은 행렬곱에 소요되는 시간복잡도이다. 

$T(B) = T(B / 2) + O(N^3)$

 

즉, 해당 점화식의 깊이를 생각해보면 log_2(B)가 된다는 것을 알 수 있다. 최종적인 시간복잡도는 $O(N^3 \cdot log_2(B))$가 된다.

<소스 코드>

위의 접근법을 이용하여 코드로 구현한 결과이다.

mat 함수는 행렬 곱을 구하는 함수이고, mul은 행렬 $m$의 $k$제곱을 구하는 함수이다.

import sys

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

def mat(a, b):
    res = [[0 for _ in range(n)] for _ in range(n)]
    for i in range(n):
        for j in range(n):
            temp = 0
            for k in range(n):
                temp += a[i][k] * b[k][j]
            res[i][j] = temp % 1000
    return res

def mul(m: list, k: int):
    if k == 1:
        for i in range(n):
            for j in range(n):
                m[i][j] %= 1000
        return m
    else:
        x = mul(m, k // 2)
        if k % 2 == 0:
            return mat(x, x)
        else:
            return mat(mat(x, x), m)
res = mul(m, b)
for i in range(n):
    print(*res[i])