[백준 17401번] 일하는 세포 (Python3)

import sys,copy
input = sys.stdin.readline
sys.setrecursionlimit(10**6)
mod = 1000000007

def multiply(A,B):
  matrix = [[0]*N for i in range(N)]
  for i in range(N):
    for j in range(N):
      for k in range(N):
        matrix[i][j] += A[i][k]*B[k][j]
      matrix[i][j] %= mod
  return matrix

def cal(A,n):
  if not n:
    return I
  if n == 1:
    return A
  cal2 = cal(A,n//2)
  if not n%2:
    return multiply(cal2,cal2)
  return multiply(multiply(cal2,cal2),A)

T,N,D = map(int,input().split())

graph = [[[0]*N for i in range(N)] for i in range(T)]
for t in range(T):
  for _ in range(int(input())):
    a,b,c = map(int,input().split())
    graph[t][a-1][b-1] = c

I = [[0]*N for i in range(N)]
for i in range(N):
  I[i][i] = 1

Tgraph = copy.deepcopy(I)
for t in range(T):
  Tgraph = multiply(Tgraph,graph[t])

result = cal(Tgraph,D//T)
for t in range(D%T):
  result = multiply(result,graph[t])

for i in range(N):
  print(*result[i])

간단한 인접행렬 경로 문제

인접행렬을 이용해서 경로의 개수를 찾는 문제는 [백준 12850번] 본대 산책2 (Python3) (tistory.com) [백준 1533번] 길의 개수 (Python3) (tistory.com) 에서 다뤄보았기 때문에 쉽게 해결할 수 있었다.

 

T초의 cycle에 대한 인접행렬을 찾고, D//T만큼 거듭제곱 해준 뒤, D%T만큼 그래프를 곱해주면 된다. 물론 행렬의 거듭제곱을 할 때는 분할정복을 이용해야한다.

 

주의할 점은 분할정복 거듭제곱을 할 때, n=0일 때에 대해서 예외처리를 해주지 않으면 recursion error가 발생할 수 있다.

D가 T보다 작은 데이터가 있는 것으로 보이기 때문에, 이에 주의해야 한다. 나는 단위행렬을 만들어줘서 n=0인 경우 단위행렬을 출력하도록 하였다.

 

 

오늘의 교훈) 0은 항상 조심하자