Prim's Lazy Minimum Spanning Tree Algorithm in Haskell (Project Euler# 107)

Question

I have implemented Lazy version of Prim's Minimum Spanning Tree Algorithm. I want to improve the code structure, follow prevalent conventions and reduce code size. I am solving Project Euler #107.

---

Explanation:

I do as follows for MST:

• Initialize an adjacency list, visited set with vertice #0, min-priority heap/queue for edges with shortest edge from vertice #0, sum of edges in MST as 0.
• Find next edge, if it is not a crossing edge, recurse.
• Otherwise I include that edge (by updating visited and heap/queue) and add to MST sum the edge's weight.

Note: Non-existent edges have been assigned value -1.

---

Code:

import           Data.Array
import qualified Data.Heap       as Heap
import           Data.List.Split
import           Data.Maybe
import qualified Data.Set        as Set

dim = 40 :: Int

main :: IO ()
main = print . maximumSaving . toAdjacencyMatrix . map (wordsBy (==',')) . lines
    =<< readFile "txt/107.txt"
    where
        toAdjacencyMatrix mat = array ((0,0),(dim-1,dim-1))
            [ ((i,j), if val == "-"  then -1 else read val)
                | i <- [0..dim-1], j <- [0..dim-1], let val = mat !! i !! j]
        totalWeight network = sum [network ! (i,j) | i <- [0..dim-1], j <- [0..dim-1], i < j, network ! (i, j) > 0]
        maximumSaving network = totalWeight network - minimumSpanningTreeEdgeSum network

minimumSpanningTreeEdgeSum :: Array (Int, Int) Int -> Int
minimumSpanningTreeEdgeSum adj =
    minimumSpanningTree'
        adj
        (Set.singleton 0)
        (Heap.fromList [(adj!(0,a),(0, a)) | a <- [1..dim-1], adj ! (0,a) > 0])
        0
    where
        minimumSpanningTree' :: Array (Int, Int) Int -> Set.Set Int -> Heap.MinPrioHeap Int (Int, Int) -> Int -> Int
        minimumSpanningTree' adj visited queue sm = case Heap.viewHead queue of
            Nothing          -> sm
            (Just (weight, edge)) ->
                if isCrossingEdge edge then
                    let sm' = sm + weight in
                    let nextVertice = if Set.notMember (fst edge) visited then fst edge else snd edge in
                    let visited' = Set.insert nextVertice visited in
                    let newEdges = [(i, j) | a <- [0..dim-1], let i = min nextVertice a, let j = max nextVertice a, Set.notMember i visited, adj ! (i,j) > 0 ] in
                    let queue'' = foldl (flip Heap.insert) queue' $ map (\e -> (adj ! e, e)) newEdges
                    in minimumSpanningTree' adj visited' queue'' sm'
                else
                    minimumSpanningTree' adj visited queue' sm
            where
                queue' = fromJust $ Heap.viewTail queue
                isCrossingEdge edge = Set.notMember (fst edge) visited ||
                    Set.notMember (snd edge) visited

Answer 1

I would separate the looping queue logic from the rest. A general implementation of unrollM necessitates a state monad here.

import Control.Monad.State
import Control.Monad.Loops

unrollM :: (Monad m, Ord a) => (a -> m [a]) -> [a] -> m [a]
unrollM f = step . (Heap.fromList :: Ord a => [a] -> Heap.MinHeap a) where
  step queue = case Heap.view queue of
    Nothing -> pure []
    Just (x, queue') -> do
      as <- f x
      (x:) <$> step (foldl (flip Heap.insert) queue' as)

minimumSpanningTreeEdgeSum :: Array (Int, Int) Int -> Int
minimumSpanningTreeEdgeSum adj = sum $ map fst $ (`evalState` Set.singleton 0) $ unrollM minimumSpanningTree'
  [(adj!(0,a),(0, a)) | a <- [1..dim-1], adj ! (0,a) > 0] where
  minimumSpanningTree' :: (Int, (Int, Int)) -> State (Set.Set Int) [(Int, (Int, Int))]
  minimumSpanningTree' (weight, (from, to)) = firstM (gets . Set.notMember) [from, to] >>= \case
    Nothing -> pure []
    Just nextVertex -> do
      modify $ Set.insert nextVertex
      filterM (gets . Set.notMember . fst . snd) $ filter ((>0) . fst)
        [ (adj ! (i,j), (i, j))
        | a <- [0..dim-1]
        , let [i,j] = sort [nextVertex, a]
        ]