LeetCode: 743. Network Delay Time

다익스트라 알고리즘을 활용하여 푼다.

 

import heapq

class Solution:
    def networkDelayTime(self, times: List[List[int]], n: int, k: int) -> int:
        
        INF = int(1e9)
        
        graph = [[] for _ in range(n+1)]
        distance = [INF] * (n+1)
        
        for time in times:
            u = time[0]
            v = time[1]
            w = time[2]
            graph[u].append((w, v))
            
        self.dijkstra(graph, distance, k)
        
        max_dist = 0
        for i in range(1, n+1):
            max_dist = max(max_dist, distance[i])
        
        if max_dist == INF:
            return -1
        else:
            return max_dist
    
    
    def dijkstra(self, graph, distance, start):
        q = []
        distance[start] = 0
        heapq.heappush(q, (0, start))
        
        while q:
            dist_u, u = heapq.heappop(q)
            if distance[u] < dist_u:
                continue
            for item in graph[u]:
                dist_uv = item[0]
                v = item[1]
                if distance[v] > dist_u + dist_uv:
                    distance[v] = dist_u + dist_uv
                    heapq.heappush(q, (distance[v], v))