Skip to content

Instantly share code, notes, and snippets.

@simonh1000
Created November 10, 2014 11:51
Show Gist options
  • Save simonh1000/bf74f2db7bcaea980a4f to your computer and use it in GitHub Desktop.
Save simonh1000/bf74f2db7bcaea980a4f to your computer and use it in GitHub Desktop.
Dijkstra's shortest path using Haskell PSQueue
import qualified Data.ByteString.Char8 as BS
import qualified Data.Vector as V
import qualified Data.List as L
import qualified Data.IntSet as IS
import qualified Data.PSQueue as PSQ
type NodeName = Int
type Dist = Int
type Edge = (NodeName, Dist)
type Edges = [Edge]
type Results = [Edge]
type Graph = V.Vector Edges
type Explored = IS.IntSet
type PrioQueue = PSQ.PSQ NodeName Dist
main = do
gr <- getEdges "dijkstra.txt"
let
initalLengthOfUnexplored = V.length gr - 1
unexplored = IS.fromList [2.. initalLengthOfUnexplored]
-- initialise heap to 'infinity'
heap = PSQ.fromList $ map (\k -> k PSQ.:-> 8000000) [2..initalLengthOfUnexplored]
-- replaces distances for nodes reachable from start
heap' = updatePrioQueue 0 (gr V.! 1) heap
paths = mainLoop gr IS.empty heap' []
return paths
mainLoop :: Graph -> Explored -> PrioQueue -> Results -> Results
mainLoop gr exs hp res =
if PSQ.null hp then
res
else mainLoop gr (IS.insert n exs) hp'' $ (n,d) : res
--else hp''
where
Just (n PSQ.:-> d) = PSQ.findMin hp
hp' = PSQ.deleteMin hp
newEdges = gr V.! n :: Edges
-- filter out newEdges that point to explored
newCuttingEdges = filter (\(node,_) -> not $ IS.member node exs) newEdges
--newCuttingEdges = newEdges
hp'' = updatePrioQueue d newCuttingEdges hp'
-- updatePrioQueue takes minDist to newly explored node, and list of all edges from new node
updatePrioQueue :: Dist -> Edges -> PrioQueue -> PrioQueue
updatePrioQueue distNewlyExplored es heap = foldl go heap es
-- foldl :: (a -> b -> a) -> a -> [b] -> a
where
go :: PrioQueue -> Edge -> PrioQueue
go hp (e,d) = PSQ.adjust (\p -> min p $ distNewlyExplored + d) e hp
-- READ IN DATA **********************************************************************************
-- "1\t80,982\t163,8164\t170,2620\t\r"
getEdges :: String -> IO Graph
getEdges path = do
-- init removes a trailing '\r'
lines <- (map (init . BS.split '\t') . BS.lines) `fmap` BS.readFile path
-- insert [] as dummy first element of Vector
return $ V.fromList $ ([]:) $ map processLine lines
where
processLine :: [BS.ByteString] -> Edges
processLine (n:connections) = map splitter connections
where
splitter :: BS.ByteString -> Edge
splitter c = (myread x, myread y)
where
(x:y:_) = BS.split (',') c
--readInt :: ByteString -> Maybe (Int, ByteString)
myread = maybe (error "can't read Int") fst . BS.readInt
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment