[백준/Python]10830: 행렬 제곱

문제

https://www.acmicpc.net/problem/10830
백준 단계별로 풀어보기 - 분할 정복

---

풀이

NxN 행렬을 B번 곱한 값 % 1000 을 출력하는 문제다.

# 메모리: 31252KB / 시간: 44ms

from sys import stdin


# 1
input = stdin.readline
N, B = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(N)]
p = 1000

A = [[col % p for col in row] for row in A]

# 2
def make_matrix(arr1: list, arr2: list) -> list:
    tmp = [[0]*N for _ in range(N)]

    for i in range(N):
        for j in range(N):
            for k in range(N):
                tmp[i][j] += arr1[i][k] * arr2[k][j]
            tmp[i][j] %= p

    return tmp

# 3
def make_square(arr: list, n: int) -> list:
    if n == 1:
        return arr

    if n % 2 == 0:
        half = make_square(arr, n//2)
        return make_matrix(half, half)
    else:
        return make_matrix(arr, make_square(arr, n-1))

# 4
ret = make_square(A, B)

for line in ret:
    print(" ".join(str(l) for l in line))

 

 

단계별 풀이

1. 입력값 할당

from sys import stdin


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

A = [[col % p for col in row] for row in A]

• input을 stdin.readline으로 재정의해주고, 행(열)의 갯수 N과 곱할 횟수 B를 할당한다.
• 행렬 A를 리스트 컴프리헨션으로 저장한다. 나눠야 할 값 1000도 p로 할당해준다.
  • ❗ 여기서 중요한게, 우리는 이 입력받은 A의 원소들도 1000으로 나누어줘야한다.
• A의 각 원소들을 p로 나눈 나머지값을 다시 A에 할당한다.

 

2. 행렬곱 함수 작성

def make_matrix(arr1: list, arr2: list) -> list:
    tmp = [[0]*N for _ in range(N)]

    for i in range(N):
        for j in range(N):
            for k in range(N):
                tmp[i][j] += arr1[i][k] * arr2[k][j]
            tmp[i][j] %= p

    return tmp

• 행렬 1, 2의 곱셈값을 반환해줄 함수 make_matrix(행렬1, 행렬2)를 작성한다.
• tmp: 행렬곱을 저장할 리스트.
• 여기서 삼중 for문을 사용했다.
  • i: 행
  • j: 열
  • k: tmp[i][j]값을 계산하기 위해 N만큼 반복. 행렬1(NxN) * 행렬2(NxN) 이므로 공통되는 반복횟수도 N이다.
• 제일 안쪽의 for문 반복을 마치면 tmp[i][j]의 값이 구해지고, 이걸 p로 나눈 나머지값을 tmp[i][j]에 재할당한다.
• 모든 반복을 마치면 리스트 tmp를 반환한다.

 

3. 제곱 처리 함수 작성

def make_square(arr: list, n: int) -> list:
    if n == 1:
        return arr

    if n % 2 == 0:
        half = make_square(arr, n//2)
        return make_matrix(half, half)
    else:
        return make_matrix(arr, make_square(arr, n-1))

 

$ A^{8} $

$ = A^{4} * A^{4} $

$ = A^{2}*A^{2} * A^{2}*A^{2} ... $이다.

• B의 값에 따라 분할해줄 함수 make_square(리스트, 제곱값)을 작성한다.
• 만약 n이 1이라면 그대로 리스트 arr(=A)을 반환해준다.
• n이 짝수라면,
  • make_square(arr, n//2)의 값을 half에 할당 후, 행렬곱 함수 make_matrix(half, half)를 반환한다.
• 홀수라면,
  • A^5 = A x A^2 x A^2 이므로,
  • make_matrix함수에 arr(=A), 제곱수-1을 한 값을 넣어 반환한다.

 

4. 최종 결과값 출력

ret = make_square(A, B)

for line in ret:
    print(" ".join(str(l) for l in line))

• make_square(A, B)을 실행한 결과 ret을 한 행씩 순회한다.
• 각 행의 원소를 str로 변환 후 공백" "과 함께 합쳐준 뒤 출력.