41. Cheapest Flights Within K Stops

리트코드

• 시작점에서 도착점까지의 가장 저렴한 가격을 계산하되, K개의 경유지 이내에 도착하는 가격을 리턴하라.

  • 경로가 존재하지 않을 경우 -1을 리턴한다.

1. 다익스트라 알고리즘 응용

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, K: int) -> int:
        graph = collections.defaultdict(list)
        
        #그래프 인접 리스트 구성
        for u, v, w in flights:
            graph[u].append((v, w))
            
        #큐 변수:[(가격, 정점, 남은 가능 경유지 수)]
        Q = [(0, src, K)]
        
        #우선순위 큐 최솟값 기준으로 도착점까지 최소 비용 판별
        while Q:
            price, node, k = heapq.heappop(Q)
            if node == dst:
                return price
            if k >= 0:
                for v, w in graph[node]:
                    alt = price + w
                    heapq.heappush(Q, (alt, v, k-1))
        
        return -1
        

• 가격을 시간이라고 가정한다면 최단 시간을 계산하는 경로는 앞서 다익스트라 알고리즘으로 동일하게 구현할 수 있다.

  • 다만, 여기에는 제약사항이 추가되었는데 K개의 경유지 이내에 도착해야 한다는 점이다.

  • 우선순위 큐에 추가할 때 K 이내일 때만 경로를 추가하여 K를 넘어서는 경로는 더 이상 탐색되지 않게 하면 될 것이다.

• 앞서 풀었던 다익스트라 알고리즘 코드를 살펴보자

class Solution:
    def networkDelayTime(self, times: List[List[int]], n: int, k: int) -> int:
        graph = collections.defaultdict(list)
        
        #그래프 인접 리스트 구성
        for u, v, w in times:
            graph[u].append((v, w))
            
        #큐 변수: [(소요시간, 정점)]
        Q = [(0, k)]
        dist = collections.defaultdict(int)
        
        #우선순위 큐 최솟값 기준으로 정점까지 최단 경로 삽입
        while Q:
            time, node = heapq.heappop(Q)
            if node not in dist:
                dist[node] = time
                for v, w in graph[node]:
                    alt = time + w
                    heapq.heappush(Q, (alt, v))
                    
        #모든 노드의 최단 경로 존재 여부 판별
        if len(dist) == n:
            return max(dist.values())
        return -1

• 우선 함수 명부터 수정해주고

• time이라는 변수 명도 문제에 맞게 price로 변경한다.

  • 여기서는 최저가를 계산해야 하므로

• 더 이상 전체 거리를 보관할 필요가 없기 때문에 dist 딕셔너리 삭제

  • 도착점까지의 최단 경로만 계산하면 된다.

• 전체 경로의 개수도 체크할 필요가 없기 때문에 여기서는 모두 삭제한다.

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, K: int) -> int:
        graph = collections.defaultdict(list)
        
        #그래프 인접 리스트 구성
        for u, v, w in flights:
            graph[u].append((v, w))
            
        k = 0
        Q = [(0, src, k)]
        
        #우선순위 큐 최솟값 기준으로 도착점까지 최소 비용 판별
        while Q:
            price, node, k = heapq.heappop(Q)
            if node == dst:
                return price
                
            #①   
            if k <= K:
                for v, w in graph[node]:
                    alt = price + w
                    heapq.heappush(Q, (alt, v, k))
        
        return -1
        

• 우선순위 큐에는 경유 횟수를 k로 설정하여 0부터 순서대로 함께 기입한다.

  • 이전 문제에서는 dist에 노드가 존재하는지 여부로 판별했으나 여기서는 K이내일 때만 (k <= K) 우선 순위 큐에 경로를 추가하고 K를 넘어서는 경로는 더 이상 탐색되지 않도록 ① 부분을 추가 했다.

• 계속 탐색하다가 현재 노드가 도착지라면, 결과를 리턴하고 종료한다.

  • 그러나, 큐를 끝까지 순회해도 찾지 못한다면 도착지까지 K 이내에 도달하는 경로는 존재하지 않는다는 얘기이므로 이 경우에는 -1을 리턴한다.

• k가 K를 넘어서게 된다면 더 이상 큐에 추가하는 일도 없다.

  • K값이 작을수록 빠르게 실행되고 탐색 또한 금방 종료할 것이다.

  • 변수 k와 K가 혼동되며, if k <= K:라는 비교 구문도 혼란스럽다. 이 부분을 혼동이 덜 되도록 좀 더 직관적으로 다음과 같이 개선하자.

Q = [(0, src, K)]

...

            if k >= 0:
                for v, w in graph[node]:
                    alt = price + w
                    heapq.heappush(Q, (alt, v, k - 1))

• 처음부터 입력값의 최대 경유지 값인 K를 우선순위 큐에 추가하고 경유지가 하나씩 늘 때마다 k - 1하는 형태로 변경해봤다.

• 비교구문도 if k >=0:와 같이 훨씬 더 직관적으로 처리할 수 있으며, K를 미리 선언할 필요도 없기 때문에 훨씬 더 알아보기 쉬운 간결한 코드가 됐다.