Given a graph, return the shortest path from start to end that has at most k edges

Question

Given a directed graph, return the shortest path from vertex start to end that has at most k edges

My Attempt

from heapq import heappush, heappop
def atMostK(n, edges, start, end, k):
    """
    :type edges: List[List[int]]
    :type start: int
    :type end: int
    :type k: int
    :rtype: int
    """
    graph = {}
    for (x, y, z) in edges:
        graph[x] = graph.get(x, []) + [(y, z)]
        
    table = {}
    for v in range(n):
        table[v] = (float("inf"), 0)  
    table[start] = (0, 0)
    
    
    stack = [(0,start)]
    visited = set()
    
    while(stack != []):
        node = heappop(stack)[1]
        w, cur_k = table[node]

        if(node in visited):
            continue
            
        visited.add(node)
        if(node in graph):
            for (v, weight) in graph[node]:
                
                
                cur_weight, amount = table[v]
                
                if(cur_weight > weight + w):
                    table[v] = (weight + w, amount + 1)
                    heappush(stack, (weight + w, v))
                    
    if(table[end][0] == float("inf")):
        return -1
    
    return table[end][0]
    

# fails on this

n = 3
edges = [[0,1,100],[1,2,100],[0,2,500]]
start = 0
end = 2
k = 1

print(atMostK(n, edges, start, end, k)) # 200, incorect result. Should be 500

# passes on this

n = 3
edges = [[0,1,100],[1,2,100],[0,2,500]]
start = 0
end = 2
k = 2

print(atMostK(n, edges, start, end, k)) # correct result

I tried to use dijkstra algorithm however I run into trouble whenever I need to backtrack to a previous visited node due to exceeding the k edge requirement.

So for example using this graph from start = 0, end = 2, and k = 1

it would first visit 0 then go to the 0->1 edge since its cheaper than the 0->2 edge, but inorder to go to 2 I would need another edge but im only permitted k=1 so the correct result should be 500.

It would work for graphs that dont require me to backtrack to already visited nodes but other than it wont work. How would I do this?

Answer 1

Using a slight modification of Bellman-Ford, you can solve this in O(K*(|V| + |E|)) time in the worst case, with many possible optimizations for the average case.

Run an outer loop with k iterations; on the ith iteration, we're computing the min-weight paths from start with at most i edges. You do need to do some extra copying from the usual Bellman-Ford to make sure that updates are only based on previous loop iterations.

def atMostK(n, edges, start, end, k):
    """
    :type edges: List[List[int]]
    :type start: int
    :type end: int
    :type k: int
    :rtype: int
    """

    table = {}
    for v in range(n):
        table[v] = (float("inf"), 0)
    table[start] = (0, 0)

    for _ in range(min(k, n-1)):
        new_table = table.copy()
        for u, v, weight in edges:
            if table[u][0] + weight < new_table[v][0]:
                new_table[v] = (table[u][0] + weight, u)
        table = new_table

    if table[end][0] == float("inf"):
        return -1

    return table[end][0]

You can modify Dijkstra's to do this as well, but the time complexity (and performance) is probably worse and harder to analyze. Your min-priority queue could be ordered based on (weight, num_edges, vertex), with checks to ensure that num_edges stays below k.