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

풀이

분할 정복을 이용한 거듭제곱 문제는 몇 개 풀어봤는데, 행렬을 거듭제곱하는 문제는 처음이었기에 겁이 나긴 했다. 일단 행렬의 곱셈하는 방법은 a(i, j) 요소에 대하여 A*B행렬일 때 A(i, k) * B(k, j) (k는 0부터 A의 열의 개수이다) 이기 때문에 일단 행렬을 곱하는 함수를 구현해냈다. 또한 매번 곱셈이 끝나면 수가 무한으로 커질 수 있기 때문에 1000으로 나눈 나머지 값을 저장하였다. 그 밑으로는 분할 정복을 통해 거듭제곱하는 함수를 구현해냈다. A**n일때 n을 1이 될 때까지 반으로 나눠서 A**(n//2)*A**(n//2)이런식으로 1이 되면 입력으로 주어진 A행렬을 return해주고 그 뒤로 재귀를 통해 제곱을 구하는 방식이다.

import sys
input = lambda: sys.stdin.readline().rstrip()
n, B = map(int, input().split())
A = [list(map(int, input().split()))for _ in range(n)]
def multiply(A, B):
  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]
      C[i][j] %= 1000
def div_sqrt(a, b):
  if b == 1:
    return a
  k = div_sqrt(a, b//2)
  tmp = multiply(k, k)
  if b % 2 == 0:
    return tmp
  else:
    return multiply(a, tmp)
print(div_sqrt(A, B))