Bonus problems in Haskell, 84

2015 Nov 13

Today, I went back on problem 84, and tried to implement Kruskal’s algorithm for minimum spanning tree, just for kicks. The main difference with Prim’s algorithm is that Prim only works on connected graphs, while Kruskal’s (and Boruvka’s) work on disconnected graphs too.

First we’ll define a new type for our convenience:

data ForestW a = ForestW [[a]] [(a, a, Int)]
    deriving (Show, Eq, Ord)

Where’s the convenience? Well, Kruskal’s algorithm begins with a set F of trees, which initially contains a separate tree for each vertex in the graph - we will represent F as a list of lists. Then, it merges these trees as long as they are connected by the minimum weight edge, until all edges are exhausted. ForestW will be the square hole for the square peg that is F, which we can avoid (as you’ll see) by flattening F if we wish.

-- Sort the edges in ascending order of weight,
-- and initialise F to one list (of vertices) for each vertex:
kruskal (GraphW vs es) = kruskal' [[v] | v <- vs] [] (sortBy (comparing (\(_, _, w) -> w)) es)

-- When the edges are exhausted, we have our MST (or MSFs):
kruskal' forest acc [] = ForestW forest acc

kruskal' forest   -- ^The initial forest F is one tree for each vertex of the graph
         acc      -- ^The minimum spanning forest
         (e : es) -- ^All the edges in the graph
         = kruskal' forest' acc' es
        -- We always remove the minimum weight edge of the graph:
        edge @ (a, b, w) = e

        -- Expand the MST if the minimum weight edge connects two forests:
        acc' = if forest == forest' then acc else e:acc

        -- Rearrange the forest:
        --  Find the trees containing the vertices of `edge`,
        --  If they're disjoint, concatenate them:
        forest' = if fa /= fb then fab : ((forest \\ fa) \\ fb) else forest

        -- Note that fa and fb are guaranteed to contain a single element:
        fa = filter (\x -> a `elem` x) forest
        fb = filter (\x -> b `elem` x) forest
        fab = nub $ concat (fa ++ fb)

Clumsily written, but helps think about types, and fun! I think I might spend a little more time on it to write a foldl version. Notice that Kruskal’s algorithm is slower than Prim’s if you don’t have an efficient disjoint set data structure. Have a great weekend!

« Past Future »