Floyd-Warshall Algorithm in Haskell

The algorithm has a regular access pattern so we can avoid a lot of indexing and still write it with lists, with (I think) the same asymptotic performance as the imperative version.

If you do want to use arrays for more speed, you might still want to do something similar to this with bulk operations on rows and columns rather than reading and writing individual cells.

-- Let's have a type for weights.  We could use Maybe but the ordering
-- behaviour is wrong - when there's no weight it should be like
-- +infinity.
data Weight = Weight Int | None deriving (Eq, Ord, Show)

addWeights :: Weight -> Weight -> Weight
addWeights (Weight x) (Weight y) = Weight (x + y)
addWeights _ _ = None

-- the main function just steps the matrix a number of times equal to
-- the node count.  Also pass along k at each step.
floydwarshall :: [[Weight]] -> [[Weight]]
floydwarshall m = snd (iterate step (0, m) !! length m)

-- step takes k and the matrix for k, returns k+1 and the matrix for
-- k+1.
step :: (Int, [[Weight]]) -> (Int, [[Weight]])
step (k, m) = (k + 1, zipWith (stepRow ktojs) istok m)
  where
    ktojs = m !! k  -- current k to each j
    istok = transpose m !! k  -- each i to current k

-- Make shortest paths from one i to all j.
-- We need the shortest paths from the current k to all j
-- and the shortest path from this i to the current k
-- and the shortest paths from this i to all j
stepRow :: [Weight] -> Weight -> [Weight] -> [Weight]
stepRow ktojs itok itojs = zipWith stepOne itojs ktojs
  where
    stepOne itoj ktoj = itoj `min` (itok `addWeights` ktoj)

-- example from wikipedia for testing
test :: [[Weight]]
test = [[Weight 0, None, Weight (-2), None],
        [Weight 4, Weight 0, Weight 3, None],
        [None, None, Weight 0, Weight 2],
        [None, Weight (-1), None, Weight 0]]

I don't know how to achieve peak performance, but I can give you some tips on making your code abstract so that you can work on performance tuning more easily.

First of all, it would be nice if, when you change around your data types, you don't have to rewrite everything. Right now, you've made everything concretely about lists of lists, so let's see if we can abstract that out. First, we have to figure out what your minimal matrix interface is. Glancing at your code, you appear to have initMatrix, list2matrix, rows, elemAt, and updateAtM. These are the functions that query or modify your matrix, and these are what you would need to implement to make a new version of this code for a different Matrix type.

One way to organize this interface is to make a class out of it. For instance:

class Matrix m where
  list2matrix :: [[a]] -> m a
  matrix2List :: m a -> [[a]]
  rows :: m a -> Int
  elemAt :: Int -> Int -> m a -> a
  updateAtM :: a -> (Int, Int) -> m a -> m a
  setDiag :: a -> m a -> m a

(I went ahead and added a matrix2List function for extracting your result and renamed/modified initMatrix into setDiag, which felt a little more general.)

We can then update your code to use this new class:

floydwarshall :: Matrix m => [[Weight]] -> m Weight
floydwarshall lst = fwAlg 1 $ initMatrix $ list2matrix lst

fwAlg :: Matrix m => Int -> m Weight -> m Weight
fwAlg k m | k < rows m = let n = rows m
                             m' = foldl (\m (i,j) -> updateDist i j k m) m [(i,j) | i <- [0..n-1], j <- [0..n-1]]
                        in fwAlg (k+1) m'
          | otherwise = m

initMatrix :: Matrix m => m Weight -> m Weight
initMatrix = setDiag 0

updateDist :: Matrix m => Int -> Int -> Int -> m Weight -> m Weight
updateDist i j k m =
    let w = min (elemAt i j m) (elemAt i k m + elemAt k j m)
    in updateAtM w (i, j) m

dist :: Matrix m => Int -> Int -> Int -> m Weight -> Weight
dist i j 0 m = elemAt i j m
dist i j k m = min (dist i j (k-1) m) (dist i k (k-1) m + dist k j (k-1) m)

Now all we need to do is start defining some Matrix types and see how performance is!

Let's start with lists, since you've already done this work. We'll have to use a newtype wrapper to make GHC happy, but ignoring the wrapping and unwrapping, this is morally the same as the code you wrote:

newtype ListMatrix a = ListMatrix { getListMatrix :: [[a]] }

instance Matrix ListMatrix where
  list2matrix = ListMatrix
  matrix2List = getListMatrix
  rows = length . getListMatrix
  elemAt i j (ListMatrix m) = m !! i !! j
  updateAtM a (i,j) (ListMatrix m) =
    let (firstRows, row:laterRows) = splitAt i m
        (firstCols, _:laterCols) = splitAt j row
    in ListMatrix $ firstRows <> ((firstCols <> (a:laterCols)):laterRows)
  setDiag x = go 0
    where go n m = if n == rows m then m else go (n+1) $ updateAtM x (n,n) m

(Also, I filled in elemAt and updateAtM.) You should be able to run

matrix2List @ListMatrix $ floydwarshall myList

and get the same result (and performance) that you currently have.

Now, on to the experimentation! All that's necessary is for us to define new instances of Matrix and see what happens. Perhaps we should try pure functions:

data FunMatrix a = FunMatrix { size :: Int, getFunMatrix :: Int -> Int -> a }

instance Matrix FunMatrix where
  list2matrix l = FunMatrix (length l) (\i j -> l !! i !! j)
  matrix2List (FunMatrix s f) = (\i -> f i <$> [0..s-1]) <$> [0..s-1]
  rows = size
  elemAt i j m = getFunMatrix m i j
  updateAtM a (i,j) (FunMatrix s f) = FunMatrix s (\i' j' -> if i==i' && j==j' then a else f i' j')
  setDiag x (FunMatrix s f) = FunMatrix s (\i j -> if i==j then x else f i j)

How does that perform? One problem is that the starting lookup function is still just indexing into the list of lists, which is slow. One fix would be to convert to an array or vector first and then index. Because we've nicely abstracted everything, all that would need to change is the definition of list2matrix right here, and you'll probably get a nice performance boost!

On the topic of performance, there's one other note I can point out. The definition of dist does some serious "dynamic programming". This could work fine if you were writing and reading directly into an array, but in this recursive form, you may end up doing a lot of duplicate work. One fix is to memoize. My goto memoization package is MemoTrie, which makes it really easy to memoize things. In this case, you could change dist to:

dist :: Matrix m => m Weight -> Int -> Int -> Int -> Weight
dist m = go'
  where
    go' = memo3 go
    go i j 0 = elemAt i j m
    go i j k = min (go' i j (k-1)) (go' i k (k-1) + go' k j (k-1))

That might give you a bit of a boost!

You might consider taking @Chi's advice and use STUArray, but you'll run into a problem: the STUArray interface demands that array lookups are in a monad. It's still possible to use the abstraction method I show off above, but you'll have to change the types of the functions. And, because you change the types in the interface, you'll need to update your algorithm code to be monadic. It can be a bit of a pain, but it might be necessary to get optimal performance.