[백준] / [Python] / [10830] 행렬 제곱

https://www.acmicpc.net/problem/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 = [[*map(int, input().split())] for _ in range(N)]

def power(A, B):
    N = len(A)
    answer = [[0] * N for _ in range(N)]
    
    for r in range(N):
        for c in range(N):
            s = 0
            for a, b in zip(A[r], B[:]):
                s += a * b[c]
            answer[r][c] = s % 1000
     
    return answer

def div_power(A, B):
    if B == 1:
        N = len(A)
        for x in range(N):
            for y in range(N):
                A[x][y] %= 1000
        return A
    
    T = div_power(A, B//2)
    if B % 2:
        return power(power(T, T), A)
    else:
        return power(T, T)

answer = div_power(A, B)

for a in answer:
    print(*a)