import Data.Heap as H | |
import Data.Heap (viewHead,insert) | |
import Prelude hiding (id) | |
import Control.Monad.ST (runST) | |
import Data.STRef (newSTRef,readSTRef,modifySTRef) | |
import Data.IntSet (IntSet,singleton,empty,fromList) | |
import Data.List (elem,delete) | |
type Vertex = Int | |
type Weight = Int | |
-- Assume complete graph | |
prim :: [a] -> (a -> Vertex) -> (a -> a -> Weight) -> [(Vertex,Vertex)] | |
prim xs id w = runST $ do | |
minh <- newSTRef $ H.fromList $ weights start | |
build minh verts [] | |
where | |
verts = map id xs | |
start = head verts | |
weights a = [(w a x, (a, x)) | x <- xs] | |
build mh [] es = return es | |
build mh vs es = do | |
h <- readSTRef mh | |
case viewHead h of | |
Just (_,(fromE,toE)) -> | |
if toE `elem` vs then | |
modifySTRef mh (insertList $ weights toE) | |
build (delete toE vs) ((fromE,toE):es) | |
else | |
build vs ((fromE,toE):es) | |
Nothing -> | |
modifySTRef mh (insertList $ weights $ head vs) | |
build vs es | |
insertList xs h = foldl (flip insert) h xs |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment