P379, 41번 K 경유지 내 가장 저렴한 항공권 문제 관련 질문

[GunBros]

문제점

책에 있는 풀이를 제출했을 때 TLE가 발생합니다.

제출 코드

• 문제링크

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

제출 결과

사용한 테스트 케이스

leetcode에서 채점을 위해 사용되어지는 테케입니다.

13
[[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]]
10
1
10

제가 생각하는 문제의 원인

• 이 풀이는 다익스트라를 살짝 변형했지만 사실 완전 탐색에 가까운 풀이입니다. 다익스트라 알고리즘은 heap에서 꺼낸 노드가 이미 계산이 끝난 노드면 더 이상 방문하지 않지만 이 풀이는 dst를 만나지 못하고 k가 0보다 크기만 하면 방문하므로 중복 탐색이 발생합니다.

• k가 9이고 n이 최대인 100이라고 가정했을 때 src를 포함한 50개의 노드는 모두 서로 간선으로 연결되어 있고 dst를 포함한 나머지 노드들은 나가는 간선도 들어오는 간선도 없다고 가정하면 이루어지는 탐색의 수는 대략 49^10이 됩니다.

궁금한 점

이 풀이를 어떻게 개선해야할지 궁금합니다.

[GunBros]

다음과 같은 코드로 해결했습니다!

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번째로 방문했는지 여부를 판단하기 위한 dist
        dist = []
        for _ in range(K + 2):
            dist.append(collections.defaultdict(int))
            
        # 큐 변수: [(가격, 정점, 남은 가능 경유지 수)]
        Q = [(0, src, K)]

        # 우선 순위 큐 최소값 기준으로 도착점까지 최소 비용 판별
        while Q:
            price, node, k = heapq.heappop(Q)
            dist[k + 1][node] = price
            if node == dst:
                return price
            if k >= 0:
                    for v, w in graph[node]:
                        # v가 k번째로 방문한 적이 있는지 확인
                        if v not in dist[k]:
                            alt = price + w
                            heapq.heappush(Q, (alt, v, k - 1))
        return -1

[likejazz]

@GunBros 안녕하세요. 책 출간 이후에 테스트케이스가 변경되어 다른 문제 풀이가 필요한 경우로 보입니다. 그런데 제안해주신 두 번째 풀이도 저는 타임아웃으로 풀리지가 않네요. 이 부분은 풀이가 가능한 코드를 구현하여 추후에 이 곳에 업데이트 하겠습니다.

[pythaac]

안녕하세요. 아직 issue가 open되어 있어서 작성해봅니다!
Runtime은 100ms ~ 108ms 정도 나옵니다.

기존 코드에서 추가한 내용입니다.

1. weight

• 해당 노드를 방문했던 경로의 (price, k)를 저장

2. 우선순위 큐에 push하는 조건 추가

• 해당 노드 weight의 price > push할 노드의 price
• 해당 노드 weight의 k <= push할 노드의 k

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