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

풀이

Code

import sys
input = sys.stdin.readline
def solution(n, b, array):
    # 1. 행렬 곱 함수 정의
    def multiple(array1, array2) :
        # 1-1. 출력 리스트 생성
        return_array = []
        # 1-2.
        for row in array1 :
            # 1-2-1. 리스트 생성
            row_array = []
            # 1-2-2.
            for col in zip(*array2) :
                # 원소 값 계산
                value = sum([(x * y) % 1000 for x, y in zip(row, col)]) % 1000
                # 리스트에 값 추가
                row_array.append(value)
            # 1-2-3. 출력 리스트에 값 추가
            return_array.append(row_array)
        # 1-3. 결과 반환
        return return_array
    # 2. 제곱 함수 정의
    def pow(array_pow, b) :
        # 2-1. 종료 조건 설정
        if b == 1 : return array
        # 2-2. 현재 제곱의 절반 값 불러오기
        now_array = pow(array_pow, b // 2)
        # 2-3. 분할 정복 수행
        if b % 2 == 0 : return multiple(now_array, now_array)
        else : return multiple(multiple(now_array, now_array), array)
    # 3. 초기 인자 처리
    for i in range(n) :
        for j in range(n) :
            array[i][j] %= 1000
    # 4. 제곱 함수 수행
    answer = pow(array, b)
    # 5. 결과 출력
    for line in answer :
        print(*line)

if __name__ == "__main__" :
    n, b = map(int, input().split())
    array = [list(map(int, input().split())) for _ in range(n)]
    solution(n, b, array)