[백준/12916] K-Path (Python)

Problem : https://www.acmicpc.net/problem/12916

 

Difficulty : Gold 1

 

Status : Solved

 

Time : 00:12:33

 

---

 

문제 설명

 

 

BOJ 월드에 존재하는 수많은 나라중 하나인 천나라에는 N개의 마을이 있다. 각각의 마을들 사이에는 길이 있을수도 있고 없을수도 있다. 만약 길이 있다면 그 길의 길이는 1로 동일하다.

민호는 천나라의 지도를 만들기 위해 N개의 마을들 사이의 연결성을 인접 행렬로 나타냈다. 그러다 미스테리한 이유로 길이가 K인 경로의 개수가 몇개인지 궁금해졌다.

민호가 작성한 인접행렬이 주어졌을때 길이가 K인 서로 다른 경로의 수가 몇개인지 알아보자.

 

 

입력 및 출력

 

입력

 

첫 번째 줄에 N, K (1 ≤ N ≤ 100, 1 ≤ K ≤ 109) 이 공백을 구분으로 주어진다.

다음 N개의 줄에 걸쳐 민호가 작성한 인접 행렬이 주어진다. i번 줄의 j번 수가 1이면 i번 마을과 j번 마을의 길이 있다는 얘기고 0이면 길이 존재하지 않는다는 이야기 이다.

 

출력

 

길이가 K인 경로의 수를 10^9+7로 나눈 나머지를 출력한다.

 

입력 예시

 

4 2
0 1 1 0
0 0 1 0
0 0 0 1
1 0 0 0

 

출력 예시

 

6

 

 

---

 

풀이

 

인접행렬의 특성을 생각해보자. 행렬의 K제곱을 수행했을 때, 그 행렬의 원소 mat[i][j]는 i에서 j로 k번의 경로를 거쳐 가는 경우의 수와 동일하다. 즉 행렬의 거듭제곱을 빠르게 구하는 것이 핵심이다. O(N^3logK) 시간복잡도로 분할 정복을 이용한 풀이가 가능하다.

 

 

풀이 코드

import sys
input = sys.stdin.readline

MOD = 10**9+7
N, K = map(int, input().split())
adj_mat = [list(map(int, input().split())) for _ in range(N)]

def matmul(A, B) :
  res = [[0]*N for _ in range(N)]
  for i in range(N) :
    for j in range(N) :
      for k in range(N) :
        res[i][j] = (res[i][j] + A[i][k]*B[k][j]) % MOD
  return res

def matpow(m, k) :
  if k == 1 :
    return m

  p = matpow(m, k//2)
  p = matmul(p, p)
  if k % 2 :
    p = matmul(p, m)
  return p

res = matpow(adj_mat, K)
print(sum(map(sum, res)) % MOD)

풀이 완료!