LeetCode #787 Cheapest Flights Within K Stops

읽기 전

• 불필요한 코드나 잘못 작성된 내용에 대한 지적은 언제나 환영합니다.
• 개인적으로 배운 점을 정리한 글입니다.

문제 링크

LeetCode #787 Cheapest Flights Within K Stops

문제 풀이 - 1

기본적인 우선순위 큐를 사용한 다익스트라 알고리즘을 사용하면 다음 코드와 같다.

python 코드

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])
        INF = 100000001 # 최대값으로 초기값 설정
        dist = [INF] * (n) # 거리 산출
        q = [[0, -1, src]] # 우선순위 큐
        while q:
            cost, stops, node = heapq.heappop(q)
            if node == dst: # 만약 목적지에 도착 시
                return cost # return
            if stops < K: # 아직 경유지 개수 한계가 남았다면
                stops + 1 # 경유지 체크
                for v, d in graph[node]:
                    heapq.heappush(q, [cost + d, stops, v]) # 큐에 입력
        return -1

그러나 위 코드는 n = 13, flights = [[11,12,74],[1,8,91],[4,6,13],[7,6,39],[5,12,8],[0,12,54],[8,4,32],[0,11,4],[4,0,91],[11,7,64],[6,3,88],[8,5,80],[11,10,91],[10,0,60],[8,7,92],[12,6,78],[6,2,8],[4,3,54],[3,11,76],[3,12,23],[11,6,79],[6,12,36],[2,11,100],[2,5,49],[7,0,17],[5,8,95],[3,9,98],[8,10,61],[2,12,38],[5,7,58],[9,4,37],[8,6,79],[9,0,1],[2,3,12],[7,10,7],[12,10,52],[7,2,68],[12,2,100],[6,9,53],[7,4,90],[0,5,43],[11,2,52],[11,8,50],[12,4,38],[7,9,94],[2,7,38],[3,7,88],[9,12,20],[12,0,26],[10,5,38],[12,8,50],[0,2,77],[11,0,13],[9,10,76],[2,6,67],[5,6,34],[9,7,62],[5,3,67]], src = 10, dst = 1, k = 10의 경우 TLE를 출력한다. 이유로는 cycle이 형성되기 때문에 의미없는 연산이 계속되는 것으로 보인다. 따라서 cycle 체크를 하고 불필요한 연산량을 줄여야 한다.

문제 풀이 - 2

어떤 노드를 방문하면 방문 체크를 한다. 이후 다른 경로에서 해당 노드를 재방문할 수 있으나 해당 노드를 방문하기 이전에 경유지가 더 적어야 한다.(우선순위 큐로 cost가 최소인 경로를 추출해왔기 때문)

python 코드

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
        visit = {}
        Q = [(0, src, 0)]
        while Q:
            price, node, k = heapq.heappop(Q)
            if node == dst:
                return price
            if node not in visit or visit[node] > k:
                visit[node] = k
                for v, w in graph[node]:
                    if k <= K:
                        alt = price + w
                        heapq.heappush(Q, (alt, v, k + 1))
        return -1