Prim's Lazy Minimum Spanning Tree Algorithm in Haskell (Project Euler# 107) 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
    =&lt;&lt; readFile "txt/107.txt"
    where
        toAdjacencyMatrix mat = array ((0,0),(dim-1,dim-1))
            [ ((i,j), if val == "-"  then -1 else read val)
                | i &lt;- [0..dim-1], j &lt;- [0..dim-1], let val = mat !! i !! j]
        totalWeight network = sum [network ! (i,j) | i &lt;- [0..dim-1], j &lt;- [0..dim-1], i &lt; j, network ! (i, j) &gt; 0]
        maximumSaving network = totalWeight network - minimumSpanningTreeEdgeSum network

minimumSpanningTreeEdgeSum :: Array (Int, Int) Int -&gt; Int
minimumSpanningTreeEdgeSum adj =
    minimumSpanningTree'
        adj
        (Set.singleton 0)
        (Heap.fromList [(adj!(0,a),(0, a)) | a &lt;- [1..dim-1], adj ! (0,a) &gt; 0])
        0
    where
        minimumSpanningTree' :: Array (Int, Int) Int -&gt; Set.Set Int -&gt; Heap.MinPrioHeap Int (Int, Int) -&gt; Int -&gt; Int
        minimumSpanningTree' adj visited queue sm = case Heap.viewHead queue of
            Nothing          -&gt; sm
            (Just (weight, edge)) -&gt;
                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 &lt;- [0..dim-1], let i = min nextVertice a, let j = max nextVertice a, Set.notMember i visited, adj ! (i,j) &gt; 0 ] in
                    let queue'' = foldl (flip Heap.insert) queue' $ map (\e -&gt; (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
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) =&gt; (a -&gt; m [a]) -&gt; [a] -&gt; m [a]
unrollM f = step . (Heap.fromList :: Ord a =&gt; [a] -&gt; Heap.MinHeap a) where
  step queue = case Heap.view queue of
    Nothing -&gt; pure []
    Just (x, queue') -&gt; do
      as &lt;- f x
      (x:) &lt;$&gt; step (foldl (flip Heap.insert) queue' as)

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