求解最多k次中转的最低机票价格算法超时问题优化

问题详情 #

• 存在n个城市，由m条航班连接。每条航班从城市u出发，抵达城市v，价格为w。
• 给定所有城市、航班，以及出发城市src、目的地dst，需找出最多经过k次中转的src到dst的最低价格。若无此路线，输出-1。

示例1 #

输入: 
n = 3, 航班 = [[0,1,100],[1,2,100],[0,2,500]]
出发城市src = 0, 目的地dst = 2, 最多中转次数k = 1
输出: 200

示例2 #

输入: 
n = 3, 航班 = [[0,1,100],[1,2,100],[0,2,500]]
出发城市src = 0, 目的地dst = 2, 最多中转次数k = 0
输出: 500

我的代码 #

n = 4
flights = [[0,1,100],[1,2,100],[2,0,100],[1,3,600],[2,3,200]]
src = 0
dst = 3
k = 1

def solution(n, flights, src, dst, k):
    from collections import defaultdict
    import heapq as hq
    
    graph = defaultdict(dict)
    for s, d, w in flights:
        graph[s][d] = w
    
    heap = [(0, src, k+1)]
    hq.heapify(heap)
    
    while heap:
        minCost, dest, redges = hq.heappop(heap)  # redges => 剩余可使用的航班数（对应中转次数+1）
        if redges >= 0:
            if dest == dst: 
                return minCost
            
        for neighbor, weight in graph[dest].items():
            if redges >= 0:
                hq.heappush(heap, (minCost+weight, neighbor, redges-1))
        
    return -1

遇到的问题 #

这段代码在测试小案例时能得到正确结果（比如上述测试用例输出700），但在处理以下大规模输入时出现超时（Time Limit Exceeded）：

输入:
n = 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]]
出发城市src = 10
目的地dst = 1
最多中转次数k = 10

请问是否可以在上述代码中加入break来优化超时问题？

解答 #

靠加break解决不了超时问题，你的代码用了Dijkstra算法，但没做状态剪枝，同一个城市会被反复推入堆里，尤其是大规模图里会堆出无数冗余状态，最终导致超时。

为啥break没用？ #

break只能终止当前循环，但堆里已经堆了大量无效状态，这些状态还是会被逐个弹出处理，根本问题没解决。

正确优化思路：记录每个城市的最优状态 #

我们需要维护一个状态记录数组，用来记录到达某个城市时，对应剩余中转次数的最低成本。如果再次到达同一个城市，要么剩余中转次数更少，要么成本更高，这种状态直接跳过，不用推入堆。

下面是修改后的代码（用二维数组记录状态，更严谨）：

def solution(n, flights, src, dst, k):
    from collections import defaultdict
    import heapq as hq
    
    graph = defaultdict(dict)
    for s, d, w in flights:
        graph[s][d] = w
    
    # dist[城市][剩余可飞航班数] = 到达该城市的最低成本
    # 剩余可飞航班数最多是k+1（因为k次中转对应k+1段航班）
    dist = [[float('inf')] * (k+2) for _ in range(n)]
    # 出发城市不管剩多少航班，成本都是0
    for i in range(k+2):
        dist[src][i] = 0
    
    heap = [(0, src, k+1)]
    hq.heapify(heap)
    
    while heap:
        cost, curr_city, remaining_flights = hq.heappop(heap)
        
        # 到目的地直接返回，Dijkstra堆顶是当前最小成本
        if curr_city == dst:
            return cost
        # 没航班可飞了，跳过
        if remaining_flights == 0:
            continue
        # 当前状态成本比已记录的高，没必要处理
        if cost > dist[curr_city][remaining_flights]:
            continue
        
        # 遍历所有邻接城市
        for next_city, price in graph[curr_city].items():
            new_cost = cost + price
            new_remaining = remaining_flights - 1
            # 如果新状态的成本更低，更新并推入堆
            if new_cost < dist[next_city][new_remaining]:
                dist[next_city][new_remaining] = new_cost
                hq.heappush(heap, (new_cost, next_city, new_remaining))
    
    # 找不到符合条件的路径
    return -1

为啥这样能解决超时？ #

通过记录每个城市在不同剩余航班数下的最低成本，我们直接过滤掉那些“更贵、更没潜力”的冗余状态，堆里只会保留有机会找到更优路径的有效状态，大幅减少需要处理的节点数量，自然就不会超时了。

内容的提问来源于stack exchange，提问作者meallhour