출처 : 파이썬 알고리즘 인터뷰

오늘의 문제

leetcode 787. Cheapest Flights Within K Stops

[시작점, 도착점, 가격] 을 리스트로 받아 k개의 경유지를 거쳐 목적지에 도달하는 최소 비용을 계산하는 문제이다.

생각의 흐름

최소비용 = 최단거리 => 다익스트라 알고리즘 이용

edges 리스트를 graph로 저장하기

-> 다익스트라 알고리즘에 이용되는 우선순위 큐에 탐색할 노드를 넣고

-> pop하면서 인접노드 탐색

-> k개의 경유지를 거치고 목적지에 도착하면 쌓인 가격을 반환하고

-> k개의 경유지를 거치지 못했거나, 목적지에 도착하지 않았으면 -1 반환하기

접근 1

그래프로 변환하기

graph = collections.defaultdict(list)

for u,v,w in flights:
  graph[u].append((v,w))

접근 2

다익스트라 알고리즘 구현

# Q = [(price, 출발지, 경유 수)]를 담은 우선순위 큐
Q = [(0, src, k)]

while Q:
  price, source, k = heapq.heappop(Q)
  # 인접 노드 탐색
  for v,w in graph[source]:
    if v == dst and k == 0:
      return price + w
    elif k-1 >= 0:
      heapq.heappush(Q, (price+w, v, k-1))

return -1

최종 코드

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
      graph = collections.defaultdict(list)
      
      for u,v,w in flights:
        graph[u].append((v,w))
        
      Q = [(0,src,k)]
      
      while Q:
        price, source, k = heapq.heappop(Q)
        for v,w in graph[source]:
          if v == dst and k == 0:
            return price+w
          elif k-1 >= 0:
            heapq.heappush(Q, (price+w, v, k-1))
            
      return -1