[코딩 테스트 연습(파이썬/Python)] 백준 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제곱한 결과를 출력한다.

 

[아이디어]

1) 분할 정복으로 지수를 2씩 나누어가면서 연산한다.

2) 재귀적으로 접근하되 중복 호출을 방지하기 위해 메모이제이션을 활용한다.

 

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

d = dict()
def pow_(A,B):
    def mul(A,B):
        trans_A = [x for x in zip(*B)]
        result = []
        for x in A:
            tmp = []
            for y in trans_A:
                tmp.append(sum([a*b for a,b in zip(x,y)])%1000)
            result.append(tmp)
        return result
    if B==1:
        return A
    elif B==2:
        return mul(A,A)
    elif B%2==0:
        if B not in d:
            d[B] = mul(pow_(A,B//2),pow_(A,B//2))
        return d[B]
    elif B%2:
        if B not in d:
            d[B] = mul(mul(pow_(A,B//2),pow_(A,B//2)), A)
        return d[B]

for p in pow_(A,B):
    print(*map(lambda x: x%1000, p))

처음에 마지막 출력시 1000으로 나누는 것을 하지 않아 80%에서 틀렸다고 나왔다. 문제를 다시 자세히 보니 입력 행렬에는 1000이 원소로 있을 수 있고, 이때 B가 1이면 mul 연산을 하지 않고 바로 튀어나와 문제가 생겼다. 따라서 마지막에 한 번 더 1000으로 나눈 나머지를 구해주었다.