[백준 10830] 행렬 제곱 with Python

문제 링크

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

10830번: 행렬 제곱

풀이

import sys

def prod(a, b):
    n = len(a)
    c = [[0 for __ in range(n)] for _ in range(n)]
    for i in range(n):
        for j in range(n):
            for k in range(n):
                c[i][j] += a[i][k] * b[k][j]
    remainder(c)
    return c

def remainder(A):
    n = len(A)
    for i in range(n):
        for j in range(n):
            A[i][j] %= 1000

def dnq(A, B):
    if B == 1: return A
    x = dnq(A, B//2)
    if B % 2: return prod(prod(x, x), A)
    else: return prod(x, x)

N, B = map(int, sys.stdin.readline().split())
A = []
for i in range(N):
    A.append(list(map(int, sys.stdin.readline().split())))
remainder(A)
print("\n".join(map(lambda x: " ".join(map(str, x)), dnq(A, B))))

함수 dnq에서

B가 1이면 A를 그대로 반환한다.

B가 짝수라면 A의 B제곱은 (A의 B//2제곱)*(A의 B//2제곱)과 같고,

B가 홀수라면 A의 B제곱은 (A의 B//2제곱)*(A의 B//2제곱)*A와 같다.

이 때, A의 B//2 제곱을 한 번 계산하여 변수 x에 저장해놓고 재활용하면 연산 횟수를 반으로 줄일 수 있다.

 

함수 prod는 n*n 행렬의 행렬곱을 반환하고,

함수 remainder는 행렬의 각 원소를 1000으로 나눈 나머지로 바꾼다.