[백준/Python] Gold IV #10830 행렬 제곱

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

풀이

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

def mul(q, w):
      l = len(q)
      e = [[0]*l for _ in range(n)]
      
      for i in range(n):
            for k in range(n):
                  p = 0
                  for j in range(n):
                        p += q[i][j] * w[j][k]
                  e[i][k] = p % 1000
      return e


def cal(a, b):
      if b == 1:
            for x in range(len(a)):
                  for y in range(len(a)):
                        a[x][y] %= 1000
            return a
      
      tmp = cal(a, b//2)
      if b % 2:
            return mul(mul(tmp, tmp), a)
      else:
            return mul(tmp, tmp)

ans = cal(a, b)

for i in range(n):
      for k in range(n):
            print(ans[i][k], end=" ")
      print()