787. Cheapest Flights Within K Stops

[zeikar]

Problem link

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

Problem Summary

비행 경로가 그래프로 주어질 때 k번 이하의 노드를 거쳐서 src부터 dst까지 가는 최소 비용을 구하는 문제.

Solution

쉬워 보이는데 어려운 문제.
일단 다익스트라로 풀긴 풀었는데 틀리고 시간초과도 나고... 오랜만에 다익스트라 써서 그런지 몰라도 꽤 삽질했다.

일단 최소 비용이므로 다익스트라 돌리면 되는데 문제는 k번 이하를 거쳐야 한다. 그래서 거친 횟수를 저장하기 위한 배열을 하나 추가해서 계속 갱신해주며 탐색해줘야 한다.
시간 초과를 방지하기 위해 이미 더 적은 횟수로 방문한 노드는 탐색에서 제외해줘야 한다.

코드는 https://leetcode.com/problems/cheapest-flights-within-k-stops/solutions/267200/python-dijkstra/?envType=daily-question&envId=2024-02-23 참고했다

Source Code

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        graphs = [[] for i in range(n)]
        for f in flights:            
            graphs[f[0]].append((f[1], f[2]))
        pq = [(0, src, k + 1)]
        stops = [0] * n

        while pq:
            cost, current, stop = heappop(pq)

            if current == dst:
                return cost

            if stops[current] >= stop:
                continue
            stops[current] = stop

            for next, price in graphs[current]:
                heappush(pq, (cost + price, next, stop - 1))

        return -1