Inflearn Community Q&A

안녕하세요 동찬님.

문제에서 나온 코드 구현 부분은 조금 생략한 부분이 있습니다.

그래서 노션에 적어두었는데, 혹시 이걸 말씀하신게 맞을 까

from collections import defaultdict
import heapq


class Solution(object):
    def networkDelayTime(self, times, n, k):
        graph = defaultdict(list)
        for time in times:
            graph[time[0]].append((time[2], time[1]))

        costs = {}
        pq = []
        heapq.heappush(pq, (0, k))

        while pq:
            cur_cost, cur_node = heapq.heappop(pq)
            if cur_node not in costs:
                costs[cur_node] = cur_cost
                for cost, next_node in graph[cur_node]:
                    next_cost = cur_cost + cost
                    heapq.heappush(pq, (next_cost, next_node))

        for i in range(1, n + 1):
            if i not in costs:
                return -1

        return max(costs.values())


times = [[2, 1, 1], [2, 3, 1], [3, 4, 1]]
n = 4
k = 2

s = Solution()
print(s.networkDelayTime(times, n, k))