LeetCode: 787. Cheapest Flights Within K Stops

src로부터 특정 dst까지의 최단 거리를 묻는 문제로, 다익스트라 알고리즘을 먼저 떠올릴 수 있다.

 

기본적인 형태의 다익스트라 알고리즘은, priority queue에 "최단 거리가 갱신되는 경우"를 추가한다.

 

하지만 이 문제는 "이동 횟수 제한"이 있으므로, 이를 수정해야한다.

 

즉, priority queue에 "최단 거리가 갱신되는 경우"를 추가하는 것이 아니라, "추가적으로 이동 가능한 경우"를 모두 추가할 필요가 있다.

 

이를 코드로 구현하면 아래와 같다.

import heapq


class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:

        INF = int(1e9)
        graph = [[] * n for _ in range(n)]
        for a, b, c in flights:
            graph[a].append((c, b))
        dist_list = [INF] * n

        dist_list[src] = 0
        q = []
        # (dist_u, u, num_move)
        heapq.heappush(q, (0, src, 0))

        while q:
            dist_u, u, num_move = heapq.heappop(q)

            if u == dst:
                return dist_u

            if num_move >= k + 1:
                continue

            for dist_uv, v in graph[u]:
                heapq.heappush(q, (dist_u + dist_uv, v, num_move + 1))

        return -1

 

물론 상기 코드를 그대로 제출하면 time limit으로 accept되지 않는다.

 

이는 priority queue에서 나오는 모든 경우를 다 테스트하기 때문이다.

 

그러나, 아래와 같은 케이스는 테스트가 불필요함을 알 수 있다.

 

: 해당 접점까지의 이동 횟수 >= 해당 접점까지 최단거리 경로의 이동 횟수

 

따라서 상기 조건을 만족하는 경우는 배제해준다.

 

import heapq


class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:

        INF = int(1e9)
        graph = [[] * n for _ in range(n)]
        for a, b, c in flights:
            graph[a].append((c, b))
        dist_list = [INF] * n
        num_move_list = [INF] * n

        dist_list[src] = 0
        q = []
        # (dist_u, u, num_move)
        heapq.heappush(q, (0, src, 0))

        while q:
            dist_u, u, num_move = heapq.heappop(q)

            if u == dst:
                return dist_u

            if num_move >= num_move_list[u]:
                continue

            if num_move >= k + 1:
                continue

            num_move_list[u] = num_move

            for dist_uv, v in graph[u]:
                heapq.heappush(q, (dist_u+dist_uv, v, num_move+1))

        return -1