[Python] 인접행렬의 제곱(백준 12850번: 본대 산책 2)

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

12850번: 본대 산책2

 

먼저 인접행렬을 통해 탐색의 경우의 수를 구하는 방법을 알아보자. 다음과 같은 행렬이 있다고 하자.

행렬1

이것을 제곱하면 아래처럼 된다.

행렬2

이것에 이동경로라는 의미를 붙여서 사용한다. 즉, AAAB + ABBB는 (A->A->B의 경우의 수) + (A->B->B의 경우의 수)가 되는 것이다. 또한 행렬2는 총 2번 이동했을 때의 경우의 수를 나타낸다.

여기서, 행렬2와 행렬1을 곱하면 3번 이동했을 때의 경우의 수가 나타난다. 그리고 행렬2를 제곱하면 4번 이동했을 때의 경우의 수가 도출된다.

즉, 행렬1을 N제곱하면 N번 이동했을 때의 경우의 수를 구할 수 있다.

 

이것을 기반으로 코드를 작성한다. 앞서 말하자면, 아래의 코드를 제출하면 TLE가 걸린다. D의 범위를 고려하면 당연한 것이며, 이해를 위해 쉬운 코드부터 설명하는 것이 목적이다.

 

import sys

input = sys.stdin.readline
MOD = 1000000007
SIZE = 8


def multiply_matrix(matrix1, matrix2):
    return [[sum(matrix1[i][k]*matrix2[k][j] % MOD for k in range(SIZE)) % MOD
             for j in range(SIZE)] for i in range(SIZE)]


def solution(D):
    matrix1 = [
        [0, 1, 1, 0, 0, 0, 0, 0],
        [1, 0, 1, 1, 0, 0, 0, 0],
        [1, 1, 0, 1, 1, 0, 0, 0],
        [0, 1, 1, 0, 1, 1, 0, 0],
        [0, 0, 1, 1, 0, 1, 0, 1],
        [0, 0, 0, 1, 1, 0, 1, 0],
        [0, 0, 0, 0, 0, 1, 0, 1],
        [0, 0, 0, 0, 1, 0, 1, 0]
    ]

    current_matrix = [
        [0, 1, 1, 0, 0, 0, 0, 0],
        [1, 0, 1, 1, 0, 0, 0, 0],
        [1, 1, 0, 1, 1, 0, 0, 0],
        [0, 1, 1, 0, 1, 1, 0, 0],
        [0, 0, 1, 1, 0, 1, 0, 1],
        [0, 0, 0, 1, 1, 0, 1, 0],
        [0, 0, 0, 0, 0, 1, 0, 1],
        [0, 0, 0, 0, 1, 0, 1, 0]
    ]

    for _ in range(D-1):
        current_matrix = multiply_matrix(current_matrix, matrix1)

    return current_matrix[0][0]


if __name__ == '__main__':
    D = int(input())
    print(solution(D))

앞서 언급한 '행렬1', 즉 1번의 이동에서의 경우의 수로 두 행렬을 초기화시킨다. 그리고 무식하게 for문을 돌면서, current_matrix를 계속해서 matrix1(행렬1)과 곱하여 정답을 도출해낸다.

하지만 앞서 언급하였듯이 TLE가 뜬다. 따라서 분할정복을 이용한 제곱을 통해 시간복잡도를 log 단위로 줄일 수 있다.

 

이에 대한 코드는 다음과 같다.

import sys

input = sys.stdin.readline
MOD = 1000000007
SIZE = 8


def multiply_matrix(matrix1, matrix2):
    return [[sum(matrix1[i][k]*matrix2[k][j] % MOD for k in range(SIZE)) % MOD
             for j in range(SIZE)] for i in range(SIZE)]


def solution(D):
    # dp 초기화
    dp1 = [
        [0, 1, 1, 0, 0, 0, 0, 0],
        [1, 0, 1, 1, 0, 0, 0, 0],
        [1, 1, 0, 1, 1, 0, 0, 0],
        [0, 1, 1, 0, 1, 1, 0, 0],
        [0, 0, 1, 1, 0, 1, 0, 1],
        [0, 0, 0, 1, 1, 0, 1, 0],
        [0, 0, 0, 0, 0, 1, 0, 1],
        [0, 0, 0, 0, 1, 0, 1, 0]
    ]

    dp2 = [[0]*SIZE for _ in range(SIZE)]
    for i in range(SIZE):
        dp2[i][i] = 1

    # 분할 정복
    while D > 0:
        if D % 2 != 0:
            dp2 = multiply_matrix(dp2, dp1)
            D -= 1
        dp1 = multiply_matrix(dp1, dp1)
        D //= 2

    return dp2[0][0]


if __name__ == '__main__':
    D = int(input())
    print(solution(D))

 

dp1은 짝수 및 제곱과 관련된 행렬을 기록하며, dp2는 홀수와 연관된 행렬을 기록한다. 위처럼 하면, dp2는 결국에 행렬1을 D제곱한 결과값을 갖게 된다. 이렇게 한 이유에 대해 의문을 갖는다면 분할정복 로직에 아래의 코드를 넣어본 뒤 코드를 돌려볼 것을 권한다.

 

    d1 = 1
    d2 = 0

    # 분할 정복
    while D > 0:
        p = D
        if D % 2 != 0:
            dp2 = multiply_matrix(dp2, dp1)
            d2 += d1
        dp1 = multiply_matrix(dp1, dp1)
        d1 *= 2
        D //= 2

        print(p, d1, d2)

123을 넣으면, 다음과 같은 결과가 나온다.

D, d1, d2순