[파이썬 알고리즘 인터뷰][최단경로] K 경유지 내 가장 저렴한 항공권

이 글은 "파이썬 알고리즘 인터뷰 (박상길 지음)"을 읽고 주관적으로 요약한 글입니다. 

---

출처 : https://www.onlybook.co.kr/entry/algorithm-interview

 

문제 정의

K개 경유지 내에 도착하는 최소 가격 리턴

 

책에서 구현된 코드

# Fail
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

 

기억해야할 기법

• queue에 push되는 요소를 줄이기 위한 조건을 찾을 것

 

내가 구현한 코드

import collections
import heapq, sys
from typing import List


class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, K: int) -> int:
        graph = collections.defaultdict(list)
        weight = [(sys.maxsize, K)] * n
        # 그래프 인접 리스트 구성
        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
                    if alt < weight[v][0] or k-1 >= weight[v][1]:
                        weight[v] = (alt, k-1)
                        heapq.heappush(Q, (alt, v, k - 1))
        return -1

• 해당 문제가 책이 발간된 이후 testcase 추가로 Time Limit 초과 발생
• 책에서 구현한 알고리즘에서 일부 조건을 추가
  • weight
    - 노드 탐색시 현재 누적 price, k를 저장
  • 우선순위 큐에 push 조건 추가
    - 현재 추가할 price가 weight의 price보다 작거나
    - 현재 추가할 k가 weight의 k보다 크거나 같을 때
• 새로 탐색하는 노드는 이전 탐색한 경로보다 더 optimal한 조건이어야함

• 관련 issue가 이미 github에 올라와 있어서 해당 코드로 건의

https://github.com/onlybooks/algorithm-interview/issues/104

P379, 41번 K 경유지 내 가장 저렴한 항공권 문제 관련 질문 · Issue #104 · onlybooks/algorithm-interview