[leet code] Cheapest Flights Within K Stops

Cheapest Flights Within K Stops - LeetCode

There are n cities connected by m flights. Each flight starts from city u and arrives at v with a price w.

Now given all the cities and flights, together with starting city src and the destination dst, your task is to find the cheapest price from src to dst with up to k stops. If there is no such route, output -1 .

문제 해석

출발지하고 도착점이 주어졌을 때 최대 k개의 정점들 지나서 도착점에 최단경로로 도달할 수 있는 경로를 찾는 문제이다. (중복 방문을 허용하고 경로가 존재하지 않는다면 -1을 return)

풀이

이 문제에서 중요한 것은 최대 K개의 정점을 지날 수 있다는 제한 조건이다.

Priority Queue를 활용할 것인데 Priority Queue에 담기는 원소는 (정점,비용,지난 정점의 갯수) 의 튜플 형태로 담길 것이다. 또한 기존의 다익스트라 알고리즘과는 다르게 Relax하는 과정이 따로 없는데, 이는 비용이 더 크더라도 지난 정점의 갯수가 더 작으면 비용이 더 작지만 지난 정점의 갯수가 더 많아 정점과 비교해 도착점에 도달할 가능성이 있기 때문이다.

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, K: int) -> int:
        
        edges = [[False for i in range(n)] for i in range(n)]
        num_edges = len(flights)
        costs = [float('inf') for i in range(n)]
        q = []
        
        for i in range(num_edges):
            start_vertice,end_vertice,path_cost = flights[i]
            edges[start_vertice][end_vertice]= path_cost
            if flights[i][0]==src:
                heapq.heappush(q,((path_cost,end_vertice,0)))
                costs[end_vertice] = path_cost
                
        while len(q)>0:
            current_cost,current_vertex,k = heapq.heappop(q)
            if current_vertex==dst:
                return current_cost
            if k<K:
                for next_vertex,cost in enumerate(edges[current_vertex]):
                    if cost:
                        costs[next_vertex] = current_cost+cost
                        heapq.heappush(q,(costs[next_vertex],next_vertex,k+1))
        return -1