How can I write a MST algorithm (Prim or Kruskal) in Haskell?

I can write both Prim's and Kruskal's algorithms to find a minimum spanning tree in C++ or Java, but I want to know how to implement them in Haskell with O(mlogm) or O(mlogn) (pure function开发者_JS百科al programs is better). Thanks a lot.

---

As svenningsson suggests, priority search queue is well suited for both Kruskal's and Prim's (atleast the author proclaims it in his paper.) The problem with Kruskal is that it requires that you have an O(log n) union-find algorithm. A union-find datastructure with a purely functional interface is described here, but it uses mutable state internally, and a purely functional implementation might be impossible and in fact, there are several problems where an efficient purely functional solution is not known, as discussed in this related SO question.

A non-pure alterenative is to implement union-find algorithm in the ST monad. A search on Hackage finds that the equivalence package suits our needs. Following is an implementation of Kruskal using Data.Equivalence.Monad from the equivalence package:

import Data.Equivalence.Monad
import Data.Graph as G
import Data.List(sortBy)
import Data.Map as M
import Control.Monad(filterM)
import Data.Ord(comparing)

run = runEquivM (const ()) (const $ const ())

kruskal weight graph = run $
 filterM go (sortBy (comparing weight) theEdges)
     where
      theEdges = G.edges graph
      go (u,v) = do 
        eq <- equivalent u v
        if eq then return False else
         equate u v >> return True

It can be used like this:

fromL xs = fromJust . flip M.lookup (M.fromList xs)

testWeights = fromL [((1,2),1),((2,3),4),((3,4),5),((1,4),30),((1,3),4)]
testGraph = G.buildG  (1,4) [(1,2),(2,3),(3,4),(1,4),(1,3)]
test = kruskal testWeights testGraph

and running test gives:

[(1,2),(1,3),(3,4)]

It should be noted that the running time is dependent on weights running in O(1) time, however fromL creates a weight function running in O(log(n)) time, this can be improved to O(1) time by using arrays or just keeping track of the weight in the input list, but it's not really part of the algorithm.

---

Here is a crude Kruskal implementation.

import Data.List(sort)
import Data.Set (Set, member, fromList, insert, union)

data Edge a = Edge a a Double deriving Show

instance (Eq a) => Eq (Edge a) where
  Edge x1 y1 z1 == Edge x2 y2 z2 = x1 == x2 && y1 == y2 && z1 == z2

instance Eq a => Ord (Edge a) where
  (Edge _ _ x) `compare` (Edge _ _ y) = x `compare` y

kruskal :: Ord a => [Edge a] -> [Edge a]
kruskal = fst . foldl mst ([],[]) . sort

mst :: Ord a => ([Edge a],[Set a]) -> Edge a -> ([Edge a],[Set a])
mst (es, sets) e@(Edge p q _) = step $ extract sets where
   step (rest, Nothing, Nothing) = (e : es, fromList [p,q] : rest)
   step (rest, Just ps, Nothing) = (e : es, q `insert` ps : rest)
   step (rest, Nothing, Just qs) = (e : es, p `insert` qs : rest)
   step (rest, Just ps, Just qs) | ps == qs = (es, sets) --circle
                                 | otherwise = (e : es, ps `union` qs : rest)
   extract = foldr f ([], Nothing, Nothing) where
       f s (list, setp, setq) =
            let list' = if member p s || member q s then list else s:list
                setp' = if member p s then Just s else setp
                setq' = if member q s then Just s else setq
            in (list', setp', setq') 

The first step is sorting the edges, which is O(n log n). The problem is to find a faster lookup for the vertex sets in the extract function. I couldn't find a faster solution for this, maybe someone has an idea...

[Update]

For a Scala implementation I used a map-like data-structure for (hopefully) better performance, but unfortunately it uses mutable sets, and I have no clue how to translate this into Haskell. The code is in my blog (sorry, description is in German): http://dgronau.wordpress.com/2010/11/28/nochmal-kruskal/

---

I think a priority search queue is what you're looking for. It can be implemented optimally in a functional language as demonstrated by Ralf Hinze in a paper. It seems like the paper is only available via acm's library at a cost.