Cheapest flights with at most k stops

Question

This is the Cheapest Flights Within K Stops problem from leetcode.com:

There are \$n\$ cities connected by \$m\$ flights. Each fight starts from city \$u\$ and arrives at \$v\$ with a price \$w\$.

Now given all the cities and fights, together with starting city \$src\$ and the destination \$dst\$, your task is to find the cheapest price from \$src\$ to \$dst\$ with up to \$k\$ stops. If there is no such route, output \$-1\$.

I have a naive implementation that fails on efficiency. I had a feeling this would happen, but I took this approach to implement my first ever iterative deepening depth first search (IDDFS). Since it fails on runtime, I don't know if my code is 100% valid from an implementation standpoint

My algorithm:

1. Enumerate all paths of depth less than \$K + 1\$ using IDDFS, add to priority queue by path cost
2. If the queue isn't empty, a path exists that makes at most \$K\$ stops, extract min from priority queue.

Obviously step 1 grows exponentially and isn't at all efficient.

What I would like help reviewing is whether my iterative deepening depth first search is from an implementation standpoint correct.

def findCheapestPrice(self, n, flights, src, dst, K):
        """
        :type n: int
        :type flights: List[List[int]]
        :type src: int
        :type dst: int
        :type K: int, K + 1 = max search depth, K + 2 = amount of visited nodes
        :rtype: int
        """
        from collections import defaultdict
        import heapq
        adj = defaultdict(list) # u : (v, cost)

        for u,v,c in flights:
            adj[u].append((v,c))

        paths = [] # pqueue

        def enumerateAllPaths(u, goal, visited, path, pathcost, depthLimit):
            visited[u] = True
            path.append(u)

            if u == goal:
                heapq.heappush(paths, (pathcost, len(path), path))
            elif depthLimit == 0:
                path.pop()
                visited[u] = False
                return
            else:
                for v,c in adj[u]:
                    if not visited[v]:
                        enumerateAllPaths(v, goal, visited, path, pathcost + c, depthLimit - 1)
            path.pop()
            visited[u] = False

        visited = [False for s in range(len(n))]
        path = []

        enumerateAllPaths(src, dst, visited, path, 0, K + 1)

        # (length of path, pathcost, path)
        bestcost = 2**63 - 1 # best cost decreases if a better path is found

        if paths:
            pathcost, pathlen, path = heapq.heappop(paths)
            bestcost = pathcost

        if bestcost != 2**63 - 1:
            return bestcost
        else:
            return -1

Answer 1

1. Review

1. There doesn't seem to be any iterative deepening going on. The code just implements an ordinary depth-first search for all simple paths in the graph. Iterative deepening would require calling enumerateAllPaths first with depthLimit=1, then with depthLimit=2 and so on.

2. The docstring should explain what the parameters mean.

3. Imports are best done once at the top level of a module, not every time a function is called.

4. The comment for paths says "pqueue" which I don't understand (is it an abbreviation for "priority queue"?). It would be better to be explicit, for example I would write something like "heap of (cost, length, path) tuples for paths that reach the goal".

5. The code for getting the best path out of the heap is needlessly long-winded, and has a bug (if the best path has cost \$2^{63}-1\$ it will not be found). Better to write:

   if paths:
       cost, _, _ = heapq.heappop(paths)
       return cost
   else:
       return -1
   

6. If you need a number that's larger than any possible cost, use infinity, that is, float('inf') instead of an arbitrary number.

7. The goal of the problem is to find just one path to the destination (the cheapest one), so it is wasteful to maintain a collections of all paths to the destination. Better just to remember the best path found so far.

8. There is no need to pass the variables goal and visited and path into the enumerateAllPaths function. These do not change during the search.

9. These lines of code are repeated:

   path.pop()
   visited[u] = False
   

   By testing depthLimit != 0 instead of depthLimit == 0 you could avoid this repetition.

10. Instead of implementing visited as a dictionary mapping nodes to True or False, use a set of nodes instead.

11. The path list is not needed: this is only used to find the length of the path, but that is redundant as len(path) is always equal to K + 1 - depthLimit. And since you don't care about the actual length of the path, only its relative length compared with other paths, you can drop the K + 1 (since this is the same for all paths) and just use -depthLimit.

2. Revised code

from collections import defaultdict

def cheapest_path_cost(edges, src, dst, k):
    """Return the cost of the cheapest path in the graph from src to dst
    visiting no more than k other nodes. The graph is specified as a
    sequence of tuples (u, v, c) meaning that there is an edge from u
    to v with cost c. If there is no path, return infinity.

    """
    graph = defaultdict(list)  # {u: (v, cost)}
    for u, v, c in edges:
        graph[u].append((v, c))
    infinity = float('inf')
    best = infinity, 0         # (cost, -depth) of best path found.
    visited = set()            # Set of nodes visited on current path.

    def enumerate_paths(u, cost, depth):
        # Enumerate simple paths of length up to depth, starting at u.
        nonlocal best
        visited.add(u)
        if u == dst:
            best = min((cost, -depth), best)
        elif depth:
            for v, c in graph[u]:
                if v not in visited:
                    enumerate_paths(v, cost + c, depth - 1)
        visited.remove(u)

    enumerate_paths(src, 0, k + 1)
    return best[0]

3. Testing

You write in the post:

Since it fails on runtime, I don't know if my code is 100% valid from an implementation standpoint.

This suggests to me that you have not tested your code. Is that right? You say in comments that you "already solved this problem with Bellman–Ford" in which case you could use your other solution as a test oracle.

Here's the kind of test that you might quickly implement using a test oracle. Here I'm using the compressed sparse graph routines in scipy.sparse.csgraph, but any other graph library would do.

import numpy as np
from scipy.sparse.csgraph import csgraph_from_dense, dijkstra
from unittest import TestCase

class TestCheapestPathCost(TestCase):
    def test_cheapest_path_cost(self):
        # Random graph on n nodes with given density and costs up to maxcost.
        n = 10
        density = 0.3
        maxcost = 5
        graph = np.random.randint(1, maxcost + 1, (n, n))
        graph[np.random.random((n, n)) > density] = 0

        # Convert to array of edges.
        i, j = np.mgrid[:n, :n]
        edges = np.dstack((i, j, graph))
        edges = edges[edges[..., 2] != 0]

        # Find shortest paths and compare.
        shortest = np.dstack((i, j, dijkstra(graph))).reshape(-1, 3)
        for src, dst, expected in shortest:
            self.assertEqual(cheapest_path_cost(edges, src, dst, n), expected)

4. Performance

There are a couple of easy optimizations you can make in depth-first searches.

First, prune paths that are no better than the best path so far, like this:

def enumerate_paths(u, cost, depth):
    # Enumerate simple paths of length up to depth, starting at u.
    nonlocal best
    visited.add(u)
    if u == dst:
        best = cost, -depth
    elif depth:
        for v, c in graph[u]:
            if v not in visited and (cost + c, 1 - depth) < best:
                enumerate_paths(v, cost + c, depth - 1)
    visited.remove(u)

Note that we don't need the call to min in this version of the code because we already tested that the cost was an improvement, before visiting this search node.

Second, when constructing the adjacency list representation of the graph, sort the edges at each vertex by their cost. This means that lower-cost paths are visited sooner in the search, and this makes it more likely that the pruning will be effective.