Programming Praxis - Minimum Spanning Tree: Prim’s Algorithm

In today’s Programming Praxis exercise we have to implement an algorithm for finding the minimum spanning tree of a graph. The Scheme solution weighs in at 15 lines, so let’s see if we can do better.

As usual, some imports:

import Data.List
import qualified Data.List.Key as K

Both the Scheme solution and the pseudocode algorithm on Wikipedia take both a list of vertices and a list of edges as input, but since the list of vertices is implicitly defined in the edges there’s really no point in specifying it separately. To get the starting vertex, we just take the first vertex of the first edge. Other than that, we do pretty much what the pseudocode algorithm says: check if there’s an edge with one connected and one unconnected point. If multiple exist, add the shortest. Add the unconnected vertex to the list of connected ones. Stop if there are no more edges with unconnected vertices.

prim :: (Eq a, Ord b) => [(a, a, b)] -> [(a, a, b)]
prim es = f [(\(v,_,_) -> v) $ head es] [] where
    f vs t = if null r then t else f (union vs [x,y]) (m:t) where
        r = filter (\(a,b,_) -> elem a vs /= elem b vs) es
        m@(x,y,_) = K.minimum (\(_,_,c) -> c) r

A quick test shows we get the same result as the Scheme version:

main :: IO ()
main = print $ prim [('A', 'D',  5), ('A', 'B', 7), ('B', 'D', 9)
                    ,('B', 'C',  8), ('C', 'E', 5), ('B', 'E', 7)
                    ,('D', 'E', 15), ('D', 'F', 6), ('E', 'F', 8)
                    ,('F', 'G', 11), ('E', 'G', 9)]

And with that we’ve reduced a 15-line solution to four lines. Not bad.

About these ads

Tags: algorithm, bonsai, code, Haskell, kata, minimum, praxis, prim, programming, spanning, tree