[leetcode 2024/02/26] 787. Cheapest Flights Within K stops

https://leetcode.com/problems/cheapest-flights-within-k-stops/description/

Cheapest Flights Within K Stops - LeetCode

 

1. 접근 방법

노드 접근 횟수에 따른, 값의 저장을 다르게 해주어야 겠다고 생각했다.

그래서 dist = [[INF for _ in range(k+2)] for _ in range(n)] 이라고 작성해 주었다. k + 2 를 한 이유는 처음과 끝 노드도 방문 노드로 count 해야하기 때문에 + 2를 해주었다.

 

즉, [[inf, inf, inf], [inf, inf, inf], [inf, inf, inf], [inf, inf, inf]] 이런식으로 선언이 된다.

1번 노드에서 4번 노드로 방문, stops가 1일 weight 가 40 일 때는,

 [[inf, 40, inf], [inf, inf, inf], [inf, inf, inf], [inf, inf, inf]] 이런식으로 저장이 될 것이다.

 

이후 방문을 할 때마다, stop + 1 을 해주었고, 해당되는 list 부분에다가 값이 적다면 갱신해주었다.

 

if node == dst 일 때, return을 그래도 해줘도 되는 이유는, min heap 이기 때문에 반환하는 값이 최솟값일 것이기 때문이다.

 

2. 코드

from typing import List
import heapq

INF = float('inf')

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        graph = [[] for _ in range(n)]
        for u, v, w in flights:
            graph[u].append((v, w))
        
        pq = [(0, src, 0)]
        
        dist = [[INF for _ in range(k+2)] for _ in range(n)]
        dist[src][0] = 0
        
        while pq:
            cost, node, stops = heapq.heappop(pq)
            
            if node == dst:
                return cost
       
            if stops > k:
                continue
            
            for next_node, w in graph[node]:
                next_cost = cost + w

                if next_cost < dist[next_node][stops + 1]:
                    dist[next_node][stops + 1] = next_cost
                    heapq.heappush(pq, (next_cost, next_node, stops + 1))

        return -1