[백준10830번] 행렬 제곱 / Python3

문제

크기가 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 5
1 2
3 4

예제 출력 1

69 558
337 406

예제 입력 2

3 3
1 2 3
4 5 6
7 8 9

예제 출력 2

468 576 684
62 305 548
656 34 412

예제 입력 3

5 10
1 0 0 0 1
1 0 0 0 1
1 0 0 0 1
1 0 0 0 1
1 0 0 0 1

예제 출력 3

512 0 0 0 512
512 0 0 0 512
512 0 0 0 512
512 0 0 0 512
512 0 0 0 512

나의 풀이

1629번 곱셈에서 거듭제곱을 분할정복으로 풀었었다. 이 문제도 이 로직을 그대로 2차원배열에 적용을 하면 풀 수 있다.

처음에 행렬은 곱하는 순서가 상관이 있다는 것을 생각하고는 분할 정복을 하면 안된다고 생각을 했지만 직접 계산해보니 n*n의 같은 행렬들을 계속 곱해주는거라 순서가 상관이 없다는 것을 알아채고는 위에서 말한 것처럼 2분할을 해가면서 거듭제곱을 구할 수 있었다.

2분할 하는 메서드와 행렬곱셈을 하는 메서드를 만들면되고 수가 커질 수 있기때문에 중간 중간 1000으로 나누어 주어야하는데 행렬곱셈을 해줄때 원소 한개의 계산을 끝낼때마다 1000으로 나눈 나머지를 저장해준다. 필자가 이렇게 풀고 계속 틀렸습니다가 나왔었는데 문제는 마지막 출력을 해줄때 1000으로 모듈라 연산을 안해주었다.

# 입력
2 1
1000 1000
1000 1000

# 결과
0 0
0 0

# 출력시에 %를 안했을경우
1000 1000
1000 1000

위의 입력과 같이 b가 1로 주어지는데 1000이상의 값이 주어진다면 출력을 해줄 때 1000으로 나눈 나머지가 아니므로 문제에 틀리게된다. 따라서 이러한 경우를 꼭 염두하고 마지막 출력을 할때도 % 1000을 해주면 된다.

코드

# 10830번 행렬 제곱

# 행렬 곱셈
def mul(n,matrix1,matrix2):
    result = [[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[i][j] += matrix1[i][k] * matrix2[k][j]
            result[i][j] %= 1000

    return result

# 2분할
def devide(n,b,matrix):
    if b == 1:
        return matrix
    elif b == 2:
        return mul(n,matrix,matrix)
    else:
        tmp = devide(n,b//2,matrix)
        if b%2 == 0:
            return mul(n,tmp,tmp)
        else:
            return mul(n,mul(n,tmp,tmp),matrix)

# 입력
n, b = map(int, input().split())
a = [list(map(int, input().split())) for _ in range(n)]

result = devide(n,b,a)
for row in result:
    for num in row:
        print(num%1000, end=' ')
    print()