문제
K부터 출발해 모든 노드가 신호를 받을 수 있는 시간을 계산하라. 불가능할 경우 -1을 리턴한다. 입력값 (u, v, w)는 각각 출발지, 도착지, 소요시간으로 구성되며, 전체 노드의 개수는 N으로 입력받는다.
입력
times = [[2,1,1], [2,3,1], [3,4,1]], N = 4, K = 2출력
2코드
from typing import List
import collections, heapq
class Solution:
def networkDelayTime(self, times: List[List[int]], N: int, K: int) -> int:
graph = collections.defaultdict(list)
for u, v, w in times:
graph[u].append((v, w))
Q = [(0, K)]
dist = collections.defaultdict(int)
while Q:
time, node = heapq.heappop(Q)
if node not in dist:
dist[node] = time
for v, w in graph[node]:
alt = time + w
heapq.heappush(Q, (alt, v))
if len(dist) == N:
return max(dist.values())
return -1풀이
이 문제에서는 다음과 같은 2가지 사항을 판별해야 한다.
- 모든 노드가 신호를 받는 데 걸리는 시간
- 모든 노드에 도달할 수 있는지 여부
첫 번째로 판별해야 하는, 모든 노드가 신호를 받는데 걸리는 시간이란 가장 오래 걸리는 노드까지의 최단 시간을 말하며, 이는 다익스트라 알고리즘(최단경로 알고리즘)으로 추출할 수 있다.
두 번째로 모든 노드에 도달할 수 있는지 여부다. 이는 모든 노드의 다익스트라 알고리즘 계산 값이 존재하는지 유무로 판별할 수 있다.
다익스트라 알고리즘 구현하기
다익스트라 알고리즘을 우선순위 큐를 최소 힙으로 구현한 모듈인 heapq를 사용하는 형태로 구현해본다.
그래프 구성
graph = collections.defaultdict(list)
for u, v, w in times:
graph[u].append((v, w))다음은 우선순위 큐를 위한 큐 변수를 정의한다. 큐 변수 Q는 '(소요시간, 정점)' 구조로 구성한다. 즉 시작점에서 '정점'까지의 소요 시간을 담아둘 것이다. 초깃값은 시작점 K부터이므로, 소요 시간은 0이고, 따라서(0, K)가 된다. dist는 거리를 의미한다.
Q = [(0, K)]
dist = collections.defaultdict(int)큐 순회를 시작하자마자 최솟값을 추출한다. 다음 dist에 node의 포함 여부를 체크한다. 이 값이 비어있으면 힙에 푸시하게되는 형태이다.
while Q:
time, node = heapq.heappop(Q)
if node not in dist:
dist[node] = time
for v, w in graph[node]:
alt = time + w
heapq.heappush(Q, (alt, v))이 문제의 예제 입력값은 너무 단순하기 때문에 좀 더 복잡한 입력값 [[3,1,5], [3,2,2], [2,1,2], [3,4,1], [4,5,1], [5,6,1], [6,7,1], [7,8,1], [8,1,1]] 정도로 해서 그림으로 나타내보자. 여기서 N = 8이고 K = 3이다.
위에는 while문을 차례로 실행했을 때의 도식표이다. 처음에 힙에 (0, 3)이 들어가 있고 이것을 pop 함으로써 시작된다. dist에 처음에는 아무것도 없으므로 dist[0] = 3이 되고, 인접 노드들을 for문으로 다 방문해서 최소 힙에 넣는 방식이다. 그러면 큐에 아무것도 없을때까지 while문이 돌고 거리가 최소인 것부터 pop하게된다. 그래야지 dist가 아에 없을 때 if문 안으로 들어가서 최소힙으로 선정된 시간이 dist value로 삽입하기 때문이다.
출처
파이썬 알고리즘 인터뷰 (글 : 박상길 그림 : 정진호) [책만]