[BOJ 10830] 행렬 제곱 - Python 풀이

문제

크기가 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

def mul_matrix(N, matrix1, matrix2):
    result_matrix = [[0 for _ in range(N)] for _ in range(N)]
    for i in range(N):
        for j in range(N):
            for k in range(N):
                result_matrix[i][j] += matrix1[i][k] * matrix2[k][j]
            result_matrix[i][j] %= 1000
    return result_matrix

def power(N, B, matrix):
    if B == 1:
        result_matrix = [[0 for _ in range(N)] for _ in range(N)]
        for i in range(N): 
            for j in range(N):
                result_matrix[i][j] = matrix[i][j] % 1000
        return result_matrix
    elif B == 2:
        return mul_matrix(N, matrix, matrix)
    else: 
        temp= power(N, B//2, matrix)
        if B % 2: 
            return mul_matrix(N, mul_matrix(N, temp, temp), matrix)
        else:
            return mul_matrix(N, temp, temp)

def solve(N, B, matrix):
    result_matrix = power(N, B, matrix)
    for row in result_matrix:
        for num in row:
            print(num, end=" ")
        print()

if __name__ == "__main__":
    N, B = map(int, sys.stdin.readline().split())
    matrix =  [[0 for _ in range(N)] for _ in range(N)]

    for i in range(N):
        input = list(sys.stdin.readline().split())
        for j in range(len(input)):
            matrix[i][j] = int(input[j])

    solve(N,B,matrix)

 

처음 딱 봤을 때 어 쉬운데? 했다가 된통 당했다...

일단 행렬 계산 하는 방법을 몰라서 구글링 해서 찾아봤다.

(https://blog.naver.com/somang8991/221438288105 <- 링크 참고)

 

그리고 제한시간이 1초길래 아 이것도 시간복잡도 O(logN)이어야 하는구나 하고 분할정복을 이용하여 풀었다.

그런데 시간초과가 났길래 찾아봤더니 어떤 사랑스러운 분의 친절한 설명을 보게 되었는데,

알고 보니 같은 결과값의 함수를 여러 번 호출해서 그런 것이었다. (https://www.acmicpc.net/board/view/76337 <- 링크 참고)

 

그래서 시간초과는 해결했는데 채점 도중 에러가 나길래 또 찾아봤더니

또 사랑스러운 분이 친절히 설명해주셨다. B가 1일때, 모듈러 연산을 하지 않고 그냥 return 해서 결과값이 틀렸기 때문이었다.

(https://www.acmicpc.net/board/view/58081 <- 링크 참고)

 

행렬 제곱을 사용해서 분할정복을 해야하는 문제라 살~짝 복잡했지만, 그래도 나름 할만했다.