문제

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

www.acmicpc.net/problem/10830

  • 입력

첫째 줄에 행렬의 크기 N과 B가 주어진다. (2 ≤ N ≤ 5, 1 ≤ B ≤ 100,000,000,000)

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

  • 출력

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

과정

일단 행렬의 곱셈을 잘 몰라서ㅎㅎ; 행렬의 곱셈부터 알아보았다.

예를 들어 4*2 행렬과 2*3을 곱하면 아래와 같이 4*3 행렬이 된다.

행렬

행렬의 제곱은 n*n 사이즈의 행렬을 거듭해서 곱하는 것이다.

이제 행렬을 제곱하는 방법을 알았다면 더 중요한 것은 행렬의 제곱을 어떻게 빠르게 계산할 것인가?이다.

아래는 그 방법을 간단하게 필기로 정리한 것이다. 지수를 2진수로 바꿔 생각해서 곱셈의 횟수를 확 줄여야 한다.

이미지

이것에 대한 자세한 내용은 여기에 정리가 잘 되어있으니 참고하자.

최종 코드

분명 생각한대로 구현했는데 처음에는 계속 답이 안 나와서 헤맸다. 내 처음 코드에는 두 가지 문제점이 있었는데

  1. 행렬 초기화 문제

결과를 저장할 행렬을 처음 입력받은 matrix로 초기화할 경우 이미 matrix를 한 번 곱한 것이 되어버린다. 따라서 정확한 결과가 나오려면 행렬을 단위행렬로 초기화하거나 b = b-1로 변경해야 한다.

  1. 1,000으로 나눈 나머지를 출력하기

이미 mul_matrix 함수에서 값을 나누어주었기 때문에 문제가 없을 것이라고 생각했지만, 출력할 때 한 번 더 1,000으로 나누어줘야 했다. 입력이

2 1 1000 1000 1000 1000

인 경우 출력할 때 한 번 더 1,000을 나누어주지 않으면 0이 아니라, 1000이 그대로 출력되기 때문이다.

이 두 가지 문제를 해결한 최종 코드는 아래와 같다.

def mul_matrix(a, b):
    global n
    res = [[0]*n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            for k in range(n):
                res[i][j] += a[i][k] * b[k][j]
                res[i][j] %= 1000
    return res

n, b = map(int, input().split())
m = [list(map(int, input().split())) for _ in range(n)]

# 결과값 res를 단위행렬로 초기화
res = [[0]*n for _ in range(n)]
for i in range(n):
    for j in range(n):
        if i == j:
            res[i][j] = 1
            break

while b:
    if b%2 == 1:
        res = mul_matrix(res, m)
    m = mul_matrix(m, m)
    b = b//2

for row in res:
    for x in row:
        #출력시 한 번 더 1000으로 나눠준다
        print(x%1000, end = ' ')
    print()

참고 사이트