Skip to content

Instantly share code, notes, and snippets.

@simonh1000
Created July 2, 2014 08:49
Show Gist options
  • Save simonh1000/e5938fb4beffe4855c8a to your computer and use it in GitHub Desktop.
Save simonh1000/e5938fb4beffe4855c8a to your computer and use it in GitHub Desktop.
Haskell Max Spanning Tree implementation
import qualified Data.ByteString.Char8 as BS
import qualified Data.Vector as V
import qualified Data.List as L
import qualified Data.Heap as H
import qualified Data.IntSet as S
type Node = Int
type Cost = Int
data Edge = Edge Node Cost deriving (Eq, Show)
type Edges = [Edge]
type Graph = V.Vector Edges
type Unexplored = S.IntSet
type MyHeap = H.MinPrioHeap Cost Node
main = do
edges <- getEdges "edges.txt"
let
gr = readGraph edges
unexplored = S.fromList [2..500]
initHeap = copyToHeap unexplored (gr V.! 1) H.empty
acc = mainLoop gr unexplored initHeap 0
return acc
mainLoop :: Graph -> Unexplored -> MyHeap -> Int -> Int
--mainLoop _ S.empty _ acc = acc
mainLoop gr unexs hp acc =
let
hp' = H.dropWhile (\(c,n) -> not $ S.member n unexs) hp
in case H.view hp' of
Just ((c,n), hp'') ->
mainLoop gr (S.delete n unexs) (copyToHeap unexs (gr V.! n) hp'') (c + acc)
Nothing -> acc
copyToHeap :: Unexplored -> Edges -> MyHeap -> MyHeap
copyToHeap ux es hp =
L.foldl' f hp es'
where
es' = filter (\(Edge n _) -> S.member n ux) es
f :: MyHeap -> Edge -> MyHeap
f h (Edge n c) = H.insert (c,n) h
-- Data in form of Lines of: Node Node Cost
getEdges :: String -> IO [(Int,Edge)]
getEdges path = do
lines <- (map BS.words . BS.lines) `fmap` BS.readFile path
let triples = (map . map) (maybe (error "can't read Int") fst . BS.readInt) lines
return $! [(a, Edge b c) | [a, b, c] <- triples]
readGraph :: [(Int,Edge)] -> Graph
readGraph es =
let
empty = V.replicate 501 []
rev = map (\(i, Edge o c) -> (o, Edge i c)) es
in V.accum (flip (:)) empty $ es++rev
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment