[BOJ] #10830 행렬 제곱 (Python, 분할정복)

 

 

🌱 문제

 

 

🌱 풀이

이 문제의 관건은 100,000,000,000까지 가질 수 있는 매우 큰 지수를 효율적으로 처리하는 것입니다. 행렬을 반복적으로 곱하는 대신, 행렬의 거듭제곱을 분할 정복(Divide and Conquer) 기법으로 계산합니다.

 

1) 입력 받기

N, B = map(int, input().split())
matrix = []

for _ in range(N):
    matrix.append(list(map(int, input().split())))

 

2) 행렬 제곱을 구하는 함수(Divide and Conquer)

def pow_matrix(matrix, B):
    """
    행렬 matrix를 B 제곱하는 함수
    """
    if B == 1:
        return [[element % 1000 for element in row] for row in matrix]

    if B == 2:
        return multiply_matrix(matrix, matrix)

    if B % 2 == 0:
        matrix = pow_matrix(pow_matrix(matrix, B//2), 2)

    else:
        matrix = multiply_matrix(pow_matrix(
            pow_matrix(matrix, B//2), 2), matrix)

    return matrix

지수가 매우 크기 때문에 행렬을 반복적으로 곱하면 시간초과가 발생하므로 분할 정복 기법을 이용합니다. 다음 수식으로 효율적으로 계산이 가능합니다. 

•  B가 짝수일 경우:  $A^B = A^{B/2} \times A^{B/2}$

•  B가 홀수일 경우:  $A^B = A \times A^{B-1} = A \times A^{(B-1)/2} \times A^{(B-1)/2}$

B가 짝수일 때와 홀수인 경우 각각의 수식으로 재귀적으로 들어가도록 작성하였습니다.

그리고 재귀함수의 base case가 되는 조건은 B가 1일 때외 2일 때입니다. B가 2일 때는 행렬의 거듭제곱을 계산(multiply_matrix() 함수)하여 리턴하고, B가 1일 때는 모든 원소를 1000으로 나눈 나머지로 바꾸어 그대로 리턴합니다!

 

3) 두 행렬을 곱하는 함수

def multiply_matrix(A, B):
    """
    두 행렬 A, B를 곱하는 함수
    """
    n = len(A)
    result = [[0] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            result[i][j] = sum(
                [(a * b) for a, b in zip(A[i], [b[j] for b in B])]) % 1000
    return result

 

행렬을 거듭제곱해주는 함수는 pow_matrix() 함수안에서 호출되는데, B=2인 기저 조건에 걸릴 때와 B가 홀수인 경우, $A \times A^{B-1}$를 계산할 때 호출합니다.

 

4) 출력하기

for row in result:
    print(' '.join(map(str, row)))

 

🌱 정답 코드

# !/usr/bin/python
# -*- coding: utf-8 -*-
# boj 10830 행렬제곱

import sys

input = sys.stdin.readline


def multiply_matrix(A, B):
    n = len(A)
    result = [[0] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            result[i][j] = sum(
                [(a * b) for a, b in zip(A[i], [b[j] for b in B])]) % 1000
    return result


def pow_matrix(matrix, B):
    if B == 1:
        return [[element % 1000 for element in row] for row in matrix]

    if B == 2:
        return multiply_matrix(matrix, matrix)

    if B % 2 == 0:
        matrix = pow_matrix(pow_matrix(matrix, B//2), 2)
    else:
        matrix = multiply_matrix(pow_matrix(
            pow_matrix(matrix, B//2), 2), matrix)
            
    return matrix


N, B = map(int, input().split())
matrix = []

for _ in range(N):
    matrix.append(list(map(int, input().split())))

result = pow_matrix(matrix, B)

for row in result:
    print(' '.join(map(str, row)))