리트코드 743번 Network Delay Time (python)

Problem

LeetCode - The World's Leading Online Programming Learning Platform

Solution

내 풀이

from typing import List
from collections import deque, defaultdict

class Solution:
    def networkDelayTime(self, times: List[List[int]], n: int, k: int) -> int:
        # visited[a] => time of n node receive signal
        visited_time = [-1] * n
        network = defaultdict(list)
        for start, end, time in times:
            network[start - 1].append((end - 1, time))
        dq = deque()

        visited_time[k - 1] = 0
        dq.append(k - 1)
        while dq:
            curr_node = dq.popleft()
            for end, time in network[curr_node]:
                # not visited or update min visited value
                if visited_time[end] == -1 or visited_time[curr_node] + time < visited_time[end]:
                    visited_time[end] = visited_time[curr_node] + time
                    dq.append(end)

        if -1 in visited_time:
            return -1
        else:
            return max(visited_time)

visited에 각 노드별 최소 도착시간을 적었다. visited에 적혀 있는 시간이 -1이거나 현재의 값보다 작다면 업데이트해주었다.

dijkstra 이용

from typing import List
from collections import deque, defaultdict
import heapq

class Solution:
    # use dijkstra
    def networkDelayTime(self, times: List[List[int]], n: int, k: int) -> int:
        network = defaultdict(list)
        for start, end, time in times:
            network[start].append((end, time))

        dist = defaultdict(int)

        # (time, node)
        heap = [(0, k)]

        while heap:
            curr_time, curr_node = heapq.heappop(heap)
            if dist.get(curr_node) is None:
                dist[curr_node] = curr_time
                for next_node, next_time in network[curr_node]:
                    heapq.heappush(heap, (curr_time + next_time, next_node))

        if len(dist) == n:
            return max(dist.values())
        else:
            return -1

heap을 이용하여 가장 작은 dist를 갖는 값부터 뽑아오기 때문에 dist에는 항상 최소의 값이 세팅되어 있다. 즉, dist에 값이 있다면 이미 최단 경로가 세팅되어 있는 것이다.

Reference

파이썬 알고리즘 인터뷰 40번