LeetCode 筆記 - 743. Network Delay Time

class Solution:
    def networkDelayTime(self, times: List[List[int]], n: int, k: int) -> int:

        # times[i] = (ui, vi, wi)
        # ui: source node
        # vi: target node
        # wi: time to reach from ui to vi
        # n: number of nodes
        # k: starting node

        def dijkstra(start):

            # priority_queue is a priority queue that stores (the shortest dis from start node, node)
            priority_queue = []
            heapq.heappush(priority_queue, (0, start_node))

            # dist is a dict that stores (node, the shortest dis from the current node)
            dis_table = collections.defaultdict(lambda: 20000)
            dis_table[start_node] = 0

            while priority_queue:
                # shortest_dis: the dis from start node
                # current_node: current node

                # pop the node with the smallest dis from start node
                shortest_dis, current_node = heapq.heappop(priority_queue)
                if dist[current_node] < shortest_dis:
                    # if the dis from start node is smaller than the dis from start node
                    # of the current node, then skip
                    continue

                # for each neighbor of the current node
                for adjacent_node, dis in graph[current_node]:
                    cost = shortest_dis + dis
                    if cost < dist[adjacent_node]:
                        dist[adjacent_node] = cost
                        heapq.heappush(priority_queue, (cost, adjacent_node))

            return dist

        graph = collections.defaultdict(list)
        for u_node, v_node, dis in times:
            graph[u_node].append((v_node, dis))

        result = dijkstra(k)
        return max(result.values()) if len(result) == n else -1