[python/백준] 10830번: 행렬 제곱

 

https://www.acmicpc.net/problem/10830

10830번: 행렬 제곱

 

 

⭐ 문제

 

 

크기가 N*N인 행렬 A가 주어진다. 이때, A의 B제곱을 구하는 프로그램을 작성하시오. 수가 매우 커질 수 있으니, A^B의 각 원소를 1,000으로 나눈 나머지를 출력한다.

첫째 줄에 행렬의 크기 N과 B가 주어진다. 둘째 줄부터 N개의 줄에 행렬의 각 원소가 주어진다. 행렬의 각 원소는 1,000보다 작거나 같은 자연수 또는 0이다.

첫째 줄부터 N개의 줄에 걸쳐 행렬 A를 B제곱한 결과를 출력한다.

 

 

 

⭐ 코드 분석

 

 

자연수의 거듭제곱을 분할정복으로 푸는 아이디어를 그대로 사용한다. 다만 여기선 행렬의 거듭제곱이므로, 행렬의 곱셈을 연산하는 함수를 정의해줘야 한다.

 

 

A^K 일때

- k가 홀수이면, A^(k//2) A^(k//2) A 로 분할

- k가 짝수이면, A^(k//2) * A(k//2) 로 분할

 

거듭제곱 함수를 그대로 구현 한 후

1. 행렬 곱셈 함수에서 나머지 연산을 정의해놨으므로 문제 조건을 충족하지만 지수가 1일 경우에는 그 구문을 거치지 않아서 나머지 연산이 수행되지 않는다.
2. 따라서 B == 1인 경우, A를 바로 리턴하지 말고, 나머지 연산을 수행 후 리턴해줘야 한다. 입력된 행렬 A의 원소 중에 1000이 있을 수도 있기 때문.

 

 

** 행렬 곱셈 함수

 

N, M = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(N)]

M, K = map(int, input().split())
B = [list(map(int, input().split())) for _ in range(M)]

AxB = [[0]*K for _ in range(N)]
for row in range(N):
    for col in range(K):
        e = 0
        for i in range(M):
            e += A[row][i] * B[i][col]
        AxB[row][col] = e

for r in AxB:
    print(*r)

 

 

 

⭐ 코드

 

import sys
input = sys.stdin.readline

N, B = map(int, input().split())
A = [[*map(int, input().split())] for _ in range(N)]

def mul(U, V):
    n = len(U)
    Z = [[0]*n for _ in range(n)]
    
    for row in range(n):
        for col in range(n):
            e = 0
            for i in range(n):
                e += U[row][i] * V[i][col]
            Z[row][col] = e % 1000

    return Z

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

result = square(A, B)
for r in result:
    print(*r)