Improved dijkstra algorithm in Haskell.

dijkstra.hs

{-

Command-line: ./dijkstra graph.txt startNode endNode
Haskell: dijkstra startNode endNode aGraph

Works with any directed edge-weighted graph.

Graph File Structure:

Node ConnectedNode DistanceToConnectedNode ConnectedNode DistanceToConnectedNode ...
Node ...
...

-}

import qualified Data.Map as Map
import Data.Tuple
import Data.Maybe

import System.Environment


type Node = String
type Distance = Int

type Path = (Distance, [Node])

data Way = Way Path | Visited deriving (Ord, Eq)
type Ways = Map.Map Node Way

type Connections = Map.Map Node Distance
type Graph = Map.Map Node Connections


main = do
    [graphFilePath, startNode, endNode] <- getArgs
    graphString <- readFile graphFilePath
    let graph = parseGraphString graphString
    print $ dijkstra startNode endNode graph

parseGraphString :: String -> Graph
parseGraphString s = foldl parseGraphStringLine Map.empty $ lines s

parseGraphStringLine :: Graph -> String -> Graph
parseGraphStringLine g l = Map.insert (head w) (parseConnectionsList Map.empty $ tail w) g where w = words l

parseConnectionsList :: Connections -> [String] -> Connections
parseConnectionsList c [] = c
parseConnectionsList c (n:dS:xs) = parseConnectionsList (Map.insert n (read dS) c) xs


dijkstra :: Node -> Node -> Graph -> Maybe [Node]
dijkstra s e g = fmap reverse $ find s e g ws
    where ws = Map.singleton s (Way (0,[s]))

find :: Node -> Node -> Graph -> Ways -> Maybe [Node]
find n e g ws | n == e =
    let (Just (Way (_, ns))) = Map.lookup e ws
    in  Just ns
find n e g ws = do
    connections <- Map.lookup n g
    let complementedWs = complement ws n connections
        newWs = Map.insert n Visited complementedWs
    nextN <- nearest newWs
    find nextN e g newWs

complement :: Ways -> Node -> Connections -> Ways
complement ws n connections =
    Map.foldWithKey (complementWay (fromWay $ fromJust $ Map.lookup n ws)) ws connections
    
complementWay :: Path -> Node -> Distance -> Ways -> Ways
complementWay gonePath n d ws
    | w == Nothing ||
      (fromJust w /= Visited &&
       newP < (fromWay $ fromJust w)) = Map.insert n (Way newP) ws
    where w = Map.lookup n ws
          newP = connect gonePath n d
complementWay _ _ _ ws = ws
    
connect :: Path -> Node -> Distance -> Path
connect (ds, ns) n d = (ds + d, n:ns)

nearest :: Ways -> Maybe Node
nearest ws =
    case filter (\(w,_) -> w /= Visited) $ map swap $ Map.toList ws of
        [] -> Nothing
        xs -> Just $ snd $ minimum xs
        
fromWay :: Way -> Path
fromWay (Way p) = p

graph.txt

A0 A1 50
B0 B1 10
A1 A0 50 B1 30 A2 5
B1 B0 10 A1 30 B2 90
A2 A1 5 B2 20 A3 40
B2 B1 90 A2 20 B3 2
A3 A2 40 B3 25 A4 10
B3 B2 2 A3 25
A4 A3 10
B4 A0 500