Computing All-Pairs Shortest Paths with Memoization

The Floyd–Warshall algorithm is a classic application of dynamic programming. It finds the shortest path between all pairs of nodes in a directed graph in $O(n^3)$ time. It can actually be applied to any semiring; for example, by replacing the operations min and + with or and and, it computes the transitive closure of a graph in cubic time.

What is less well known is that the algorithm can easily be expressed using memoization, with the same asymptotic performance and the benefit that the correctness of the algorithm is, arguably, more obvious.

Let $d(i,j)$ be the one-step distance from node $i$ to node $j$, or ∞ if there is no direct edge from $i$ to $j$. The goal of the algorithm is to compute, for all pairs of nodes $i$ and $j$, the minimum distance along some path from $i$ to $j$. Let's write this shortest distance as $i⤳j$. The key to the success of dynamic programming (and memoization) is that the problem has optimal substructure: if the node $k$ lies on the shortest path from $i$ to $j$, then $i⤳j=(i⤳k)+(k⤳j)$.

Assume that the $n$ graph nodes are numbered $0,\dots ,n-1$ so that we can write $i<j$ to mean that node $i$ has a lower index than $j$. Then given a best path $i⤳j$ between two nodes, there must be some index $k$ such that all other nodes along the path have an index at least $k$. To support enumerating the set of paths being considered, we define a function that computes shortest paths between two nodes while considering only the paths whose other node indices are at least some given $k$:

    (* sp(i, j, k) is the length of the shortest path from node i
     * to node j, where other than nodes i and j, the path includes only nodes
     * whose index is at least k.
     *)
    val sp: int * int * int -> int

To compute the best path distance $i⤳j$, this function is simply invoked with $k=0$: sp(i, j, 0).

Eliding memoization, this function has a straightforward recursive definition that follows from its specification:

let rec sp(i, j, k) =
  if k >= n then d(i, j) else    (* No other nodes can be on path *)
  min
    (sp(i, k, k+1) + sp(k, j, k+1)) (* Best path including node k *)
    (sp(i, j, k+1))                 (* Best path excluding node k *)

As written, this code takes exponential time. However, in a graph with $n$ nodes, there are only $n^3$ distinct arguments to this function. Therefore, once memoization is added, the function can be applied to all pairs $(i,j)$ while incurring only $O(n^3)$ evaluations of this code.

Adding memoization to this code can be done in multiple ways. In OCaml, we can define a function memoize that generically memoizes recursive functions, using a hash table:

let memoize f =
  let h = Hashtbl.create 11 in
  let rec memoized x =
    try
      Hashtbl.find h x
    with
      Not_found ->
        let y = f memoized x in
          Hashtbl.add h x y;
          y
  in
    memoized

Next, we update the function sp to take an extra argument that will be its memoized self:

let sp sp (i, j, k) =
  if k >= n then d(i, j) else
  min
    (sp(i, k, k+1) + sp(k, j, k+1))
    (sp(i, j, k+1))

And we can apply memoize to this function to obtain the efficient version of the algorithm!

let memo_sp = memoize sp
let shortest_path i j = memo_sp(i,j,0)

Andrew Myers