백준 14289 본대 산책 3

14289 본대 산책 3

문제

양방향 그래프가 주어진다.
모든 간선의 비용(소요시간)은 1이다.
자연수 D가 주어질 때 정점 1에서 출발해 정점 1로 D시간만에 돌아오는 경로의 개수를 출력한다.

풀이

아이디어가 거의 전부인 문제
유사한 문제를 풀어보았거나 인접행렬을 이용한 그래프 이동을 떠올린다면 구현은 아주 쉽다.

인접행렬의 곱셈

$i×i$형태의 인접행렬 $A$가 아래와 같이 있다.

이 때 인접행렬 $A$의 제곱은 다음과 같이 표현된다.

그런데 문제에서는 모든 비용이 1이라고 했으므로 이동할 수 없는 경우의 비용을 0으로 표현하면

이므로

비슷한 논리로 3이상의 $n$에서도 $A^n$의 $i$행 $j$열의 원소는 $i$에서 $j$까지 $n$개의 정점을 거쳐 이동하는 경우의 수이다.

따라서 문제는 주어진 인접행렬 $A$로부터 $A^D$의 $0,0$원소의 값을 구하는 문제로 귀결된다.

굳이 행렬 곱셈으로

bfs나 dfs로 찾으면 $O(N^D)$
행렬 곱을 이용하면 별도의 행렬곱 알고리즘 없이도 분할정복을 이용한 거듭제곱만으로
$O(N^3log(N))$로 계산할 수 있다.

코드

def matrix_multiply(A: list[list[int]], B: list[list[int]]):
    C = [[0]*N for _ in range(N)]
    for i in range(N):
        for j in range(i+1):
            C[i][j] = C[j][i] = sum(A[i][k] * B[j][k] for k in range(N)) % 1000000007

    return C

def matrix_power(arr: list[list[int]], n: int):
    if n in cache:
        return cache[n]
    
    a = n // 2
    b = n - a

    result = matrix_multiply(matrix_power(arr, a), matrix_power(arr, b))
    cache[n] = result
    return result

N, M = map(int, input().split())
A = [[0]*N for _ in range(N)]

for _ in range(M):
    a, b = map(int, input().split())
    A[a-1][b-1] = A[b-1][a-1] = 1

D = int(input())

cache = {1: A}

print(matrix_power(A, D)[0][0])