acatalepsie/content/posts/haskell-dijkstra.lhs.md at c08195444eb1ecdaff7a31ff2480eb0ca7818988

12 KiB

Raw Blame History

title	date	draft
Generalized Dijkstra in Haskell	2024-12-20	true

This years' Advent of Code has lots of 2D grids, and makes you do many traversals on them to find paths of various kinds. At some point I had to implement Dijkstra's algorithm, in Haskell. In trying to make my implementation reusable for the following days, I realized that Dijkstra's algorithm is actually way more general than I remembered --- or was taught! In short, weights don't have to be real-valued!

In this post, I describe a general interface for the algorithm, such that we can implement it exactly once and still use it to compute many different things.

This article is a literate Haskell file, so feel free to download it and try it for yourself! As such, let's get a few imports and language extensions out of the way:

<details>
  <summary>Haskell Bookkeeping</summary>

{-# LANGUAGE GHC2021 #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE RecordWildcards #-}
{-# LANGUAGE StandaloneKindSignatures #-}

import Data.Function ((&))
import Data.Functor ((<&>))
import Data.Kind (Type)
import Data.Set (Set)
import Data.Set qualified as Set

</details>

A primer on Dijkstra's algorithm

Let's recall the problem statement that Dijkstra's algorithm solves, and the pseudo-code of the algorithm. If you know that already, you can skip to the next section.

The shortest path problem

Consider a weighted directed graph G = (V, E, w).

V denotes the set of vertices.
E \subseteq V \times V denotes the set of edges.
Every edge e \in E has an associated (non-negative) weight w(e) \in \mathbb{R}, w(e) ≥ 0.

We call path a sequence of vertices such that there is an edge between every consecutive vertex in the sequence. If we denote \text{paths}(a, b) the set of paths from a to b in G, this means $p = \langle v_0, \dots, v_k \rangle \in \text{paths}(a, b) if $v_0 = a, v_k = b and \forall 0 ≤ i < k, (v_i, v_{i + 1}) \in E.

We can define the weight of a path as the sum of the weights of its constituent edges.

w(p) = \sum_{i = 1}^k{w(v_{i - 1}, v_i)}

Shortest-paths problems ask questions along the line of:

What is the minimum weight from a to b in G?
Can we find one path from a to b with minimum weight?
Can we find all such paths from a to b?

If you interpret the weight as a physical distance, this amounts to finding the shortest trip from one vertex to the other.

Dijkstra's algorithm is an infamous technique for solving the single-source shortest-paths problem: finding a shortest path from a given source vertex s to every vertex v \in V. It's essentially a generalization of breadth-first search to (positively) weighted graphs. And it's pretty fast!

Dijkstra's algorithm

Taking a step back

One thing to notice in the problem statement from earlier is that weights have very little to do with real numbers. In fact, they don't have to be scalars at all!

If we denote W the set of weights, the algorithm merely requires:

An equivalence relation (\cdot \approx \cdot) \subseteq W \times W on weights.
It doesn't have to be definitional equality!
A total order on weights, that is: (\cdot\leq\cdot) \subseteq W \times W such that it is transitive, reflexive and anti-symmetric. This order should be compatible with \approx i.e. equivalence preserves order.
A way to add weights together $(\cdot \oplus \cdot) ∷ W \rightarrow W \rightarrow W$, such that:
- \oplus is associative.
- \approx is compatible with \oplus.
- \leq is compatible with \oplus.
- x \oplus y is an upper bound of both x and y, i.e "adding costs together can only increase the total cost".
A neutral element 0 for \oplus, that should also be a lower bound of W.
An absorbing element \infty for \oplus, that should also be an upper bound of W.

If we summarize, it looks like (W/\approx, \oplus, 0) is a monoid, totally ordered by \leq and with null element \infty. I think this encompasses all the properties stated above, and nothing more, but I haven't looked that deeply into the formalism and how mathematicians usually call these things.

Now we can freely redefine the weight of a path:


w(p) = \bigoplus_{i = 1}^k{w(v_{i - 1}, v_i)}

Equipped with this toolkit, we can redefine the single-source shortest-path problem: for a given source vertex s \in V, how do we compute the smallest weight achievable on a path from s to any other vertex e \in V?

We can see that the pseudo-code for Dijkstra's algorithm remains almost identical:

S = \emptyset
Q = \emptyset
For each vertex u \in V:
- Insert u \in Q
While Q \neq \emptyset:
- u = \text{Extract-Min}(Q)
- S = S \cup \{u\}
- For each vertex v in such that (u, v) \in E:

Some interesting remarks:

(\cdot \oplus \cdot) does not have to be commutative!

Given the requirements on weights we established earlier, we can try to map it to the corresponding Haskell implementation.

Weights should have an equivalence relation: that's Eq.
Weights should have a total order: that's Ord.
Weights should have an associative addition operation that respects the order: that's Semigroup.

Sadly we're not using Agda so we can't enforce the fact that the order relation must be compatible with the semigroup operation from inside the language, so we'll just have to be careful when instanciating.

So, a cost should satisfy have instances for all three classes above (and Ord implies Eq in Haskell, somehow).

class (Semigroup a, Ord a) => Cost a where
  merge :: a -> a -> a
  merge x = const x

Add corresponds to semigroup The default implementation is to ignore the new cost. This is quite common, say, if we only want to find "a shortest path", and not every one of them.

A generic interface

data Dijkstra i c = Dijkstra
    { bounds      :: (i, i)
    , startCost   :: i -> c
    , defaultCost :: c
    , next        :: i -> c -> [(c, i)]
    }

So, let's expand a bit on the fields of the interface.

bounds describes the lower and upper bound of the vertices. This is only necessary because I want to store intermediate costs in a mutable array during the traversal, for efficiency purposes. If you cannot reasonnably enumerate all vertices, you can drop the bounds field and use a pure persistent Map instead.
initCost returns the initial cost we use for the start vertex.
defaultCost is the initial cost for all other (yet unvisited) vertices.

Given a vertex and its associated cost, the transition function next returns its neighbours in the graph, along with the updated cost to get there. For it to make sense, the cost of neighbours should be higher then (or equal to) the cost of the parent vertex.

Generic Dijkstra implementation

Now, let's implement this thing.

We first need a priority queue, with the following interface:

type PQueue :: Type -> Type

emptyQ     :: PQueue a
singletonQ :: Ord a => a -> PQueue a
insertQ    :: Ord a => a -> PQueue a -> PQueue a

pattern EmptyQ :: PQueue a
pattern (:<)   :: Ord a => a -> PQueue a -> PQueue a

For simplicity, let's just use a wrapper around Data.Set.

<details>
  <summary><code>PQueue</code> implementation</summary>

newtype PQueue a = PQueue (Set a)

emptyQ = PQueue Set.empty
singletonQ = PQueue . Set.singleton
insertQ x (PQueue s) = PQueue (Set.insert x s)

minView :: PQueue a -> Maybe (a, PQueue a)
minView (PQueue s) =
  case Set.minView s of
    Just (x, s') -> Just (x, PQueue s')
    Nothing      -> Nothing

pattern EmptyQ   = (Set.minView -> Nothing)
pattern (:<) x q = (Set.minView -> Just (x, q))

</details>

And last, the implementation for Dijkstra's algorithm:

dijkstra :: (Ix i, Cost c) => Dijkstra i c -> i -> Array i c
dijkstra (Dijkstra{..} :: Dijkstra i c) start = runST do
  -- create a mutable array to store the cost of all vertices
  costs :: STArray s i c <- newArray bounds defaultCost
  -- update the cost of the starting position
  let startCost = initCost start
  writeArray costs start startCost
  -- traverse the graph starting from the start
  aux costs (Set.singleton (startCost, start))
  -- return the minimal costs of all vertices
  freeze costs
  where
    aux :: STArray s i c -> PQueue (c, i) -> ST s ()

Instanciating the interface

The interface for dijkstra's algorithm is very abstract now. Let's see how to instanciate it to compute useful information!

But first, we define a basic graph that we will traverse in all the following examples.

type Coord = (Int, Int)
type Grid  = Array Coord Char

neighbours :: Coord -> [Coord]
neighbours (x, y) =
    [ (x - 1, y    )
    , (x + 1, y    )
    , (x    , y - 1)
    , (x    , y + 1)
    ]

emptyCell :: Grid -> Coord -> Bool
emptyCell grid = ('#' /=) . (grid !)

graph :: Array Coord Char

Minimum distance

The simplest example is to try and compute the length of the shortest path between two vertices. We define our cost, in this case an integer for the length.

newtype MinDist = MinDist !Int deriving (Eq, Ord, Cost)

And then we instanciate the interface.

minDist :: Grid -> Dijkstra Coord MinDist
minDist grid = Dijkstra 
  { bounds      = IArray.bounds grid
  , initCost    = const (MinDist 0)
  , defaultCost = MinDist maxBound
  , next        = next
  }
  where next :: Coord -> MinDist -> [(MinDist, Coord)]
        next x (MinDist d) =
          neighbours x
          & filter (inRange (IArray.bounds grid))
          & filter ((/= '#') . (grid !))
          & map (MinDist (d + 1),)

getMinDist :: Grid -> Coord -> Coord -> MinDist
getMinDist = dijkstraTo . minDist

Shortest path

Maybe this interface has become too abstract, so let's see how to instanciate it to find the usual shortest path. We introduce a new cost that stores a path along with its length.

data Path i = Path !Int [i]

Given that we only want to find a shortest path, we can put paths with the same length in the same equivalence class, and compare paths only by looking atChar their length.

instance Eq  (Path i) where
  Path l1 _ == Path l2 _ = l1 == l2

instance Ord (Path i) where 
  Path l1 _ `compare` Path l2 _ = compare l1 l2

Again, because we only want to find a shortest path, if we merge two paths in the same equivalence class, we simply return the first one. Lucky for us, that's the default implementation for merge in Cost (Path i).

instance Cost (Path i) where

And... that's all we need! Running it on a sample graph gives us a shortest path.

Shortest paths (plural!)

Now what if we want all the shortest paths? Simple, we define a new cost!

data Paths i = Paths !Int [[i]]

The only difference with Path is that we store a bunch of them. But we compare them in the exact same way.

instance Eq  (Path i) where
  Path l1 _ == Path l2 _ = l1 == l2

instance Ord (Path i) where 
  Path l1 _ `compare` Path l2 _ = compare l1 l2

However, when we merge costs we make sure to keep all paths:

instance Cost (Paths i) where
  merge (Paths l xs) (Paths _ ys) = Paths l (xs ++ ys)

And... that's it!

<!-- TODO: pre-compile katex to MathML only -->
<!-- <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.3/dist/katex.min.css" integrity="sha384-Juol1FqnotbkyZUT5Z7gUPjQ9gzlwCENvUZTpQBAPxtusdwFLRy382PSDx5UUJ4/" crossorigin="anonymous"> -->
<script defer src="https://cdn.jsdelivr.net/npm/katex@0.16.3/dist/katex.min.js" integrity="sha384-97gW6UIJxnlKemYavrqDHSX3SiygeOwIZhwyOKRfSaf0JWKRVj9hLASHgFTzT+0O" crossorigin="anonymous"></script>
<script type="module">
  const macros = {}
  document.querySelectorAll('.math').forEach(elem => {
    katex.render(elem.innerText, elem, {throwOnError: false, macros,
    displayMode: !(elem.classList.contains('inline')), output:'mathml'})
  })
</script>

12 KiB Raw Blame History