Bug Report for cheapest-flight-path

[peizhao1]

Bug Report for https://neetcode.io/problems/cheapest-flight-path

Additional test case:

Input: n = 5, flights = [[0,1,1],[1,2,1],[2,3,1],[3,4,1],[0,2,10], src = 0, dst = 4, k = 2

Output: 12

Without this test case, my incorrect solution would pass all the existing ones:

from collections import defaultdict
import heapq

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        graph = defaultdict(list)
        for a, b, w in flights:
            graph[a].append((b, w))
        heap = [(0, src, -1)]
        shortest = defaultdict(lambda: -1)
        while heap:
            w, node, stop = heapq.heappop(heap)
            if node == dst:
                return w
            if node not in shortest:
                shortest[node] = w
            for nei, w2 in graph[node]:
                if nei not in shortest and stop + 1 <= k:
                    heapq.heappush(heap, (w + w2, nei, stop + 1))
        return shortest[dst]