[LeetCode] 743. Network Delay Time (Python)

안녕하세요!

오늘은 Leet Code 743. Network Delay Time 문제를 풀어보겠습니다. 

1~n 으로 라벨링 되어있는 n 개의 노드가 있습니다. times 라는 배열도 있는데, times[i] = (ui, vi, wi) 에서 ui는 source node, vi는 target node, wi는 source에서 target으로 신호를 보내는 데 걸리는 시간을 의미합니다. k라는 노드에서 신호를 보낼 때, 모든 n 노드가 신호를 받는 데 걸리는 최소 시간을 출력하는 함수를 만드는 문제입니다. 이때 모든 노드가 신호를 받는 것이 불가하다면 -1을 출력해주면 됩니다. 

저는 Dijkstra 알고리즘으로 문제를 풀어보려고 합니다.

---

<Dijkstra 알고리즘>

다익스트라 알고리즘은 음의 가중치가 없는 그래프에서, 가중치 그래프의 한 노드에서 다른 모든 노드까지의 최단 경로를 찾는 알고리즘입니다. Greedy 알고리즘에 기반한 이 알고리즘은 매 단계에서 가장 비용이 적은 경로를 선택하면서 점차적으로 최적의 해를 탐색해나갑니다. 

---

먼저 각 노드를 기준으로 인접 노드까지의 최단 거리를 저장할 배열 'graph'를 만들어줍니다. 이후 노드 간 거리를 모두 무한대로 초기화해주고, 시작 노드의 거리는 0으로 설정해줍니다. 노드 간 거리를 모두 무한대로 초기화하는 이유는 최종 결과로 최단 거리를 출력해야하고, 아직 도달할 수 없는 노드를 처리하기 위함입니다. 그 다음으로는 우선순위 큐를 활용해 현재 노드와 거리를 저장하고 최단 거리의 인접 노드로 업데이트 해줍니다. 모든 노드 탐색을 마친 후 최단 거리가 기록된 distance에서 최댓값을 찾아 출력해줍니다. 도달할 수 없는 노드가 있다면 -1을 반환해줍니다. 

import heapq

class Solution(object):
    def networkDelayTime(self, times, n, k):
        """
        :type times: List[List[int]]
        :type n: int
        :type k: int
        :rtype: int
        """

        # 그래프 표현
        graph = {}
        for u, v, w in times:
            if u not in graph:
                graph[u] = []
            graph[u].append((v, w))

        # 거리 표현
        # 아직 탐색하지 않은 노드의 거리는 무한대로 표현
        # 최단 경로를 찾아야하고, 도달할 수 없는 노드를 처리하기 위해
        distance = {node: float('inf') for node in range(1, n+1)}
        distance[k] = 0 # 시작 노드

        # heapq 초기 설정 (노드, 거리)
        q = [(k, 0)]

        while q:
            node, dist = heapq.heappop(q) # 현재 노드와 거리

            # 이미 처리된 노드라면 스킵
            if dist > distance[node]:
                continue

            # 인접 노드의 거리 업데이트
            if node in graph:
                for neighbor, weight in graph[node]:
                    new_dist = dist + weight

                    # 기존 거리보다 짧으면 대체
                    if new_dist < distance[neighbor]:
                        distance[neighbor] = new_dist
                        heapq.heappush(q, (neighbor, new_dist))
        
        # 모든 노드에 도달하는 데 필요한 시간
        max_dist = max(distance.values())

        # 무한대보다 작은 값이면 출력 아니면 -1 출력
        return max_dist if max_dist < float('inf') else -1