Dijkstra

Author: mnvr

17.graph.hs

@@ -3,7 +3,7 @@
 
 import Data.Map qualified as M
 import Data.Set qualified as S
-import Data.Maybe (fromJust)
+import Data.Maybe (fromJust, fromMaybe)
 import Data.Sequence (Seq(..), fromList, (><))
 
 main :: IO ()
@@ -14,7 +14,7 @@ main = interact $ unlines . demo . parse
     dbfs grid = bfs grid (0, 0) (visitor "bfs")
     dsp grid = let end = maxNode grid
                    (r, zs) = dijkstra grid (0, 0) end (visitor "shortest-path")
-               in zs ++ ["shortest-path result " ++ show r]
+               in zs ++ ["shortest-path result " ++ (show $ fromMaybe (-1) r)]
 
 type Node = (Int, Int)
 data Grid a = Grid { items :: M.Map Node a, maxNode :: Node } deriving Show
@@ -52,6 +52,30 @@ bfs grid@Grid { items } start visitor = go (Empty :|> start) S.empty
     go (xs :|> x) seen = visit x : go ys (S.insert x seen)
       where ys = (fromList $ neighbours grid x) >< xs
 
+-- For our example, the weight of the edge between u and v is the value
+-- at v. But this can be any arbitrary function or input data.
+distance :: Grid Int -> Node -> Node -> Int
+distance Grid { items } u v = fromJust $ M.lookup v items
+
 -- Find the shortest path from start, to end, using Dijkstra's algorithm.
-dijkstra :: Grid a -> Node -> Node -> (Node -> a -> b) -> (Int, [b])
-dijkstra grid@Grid { items } start end visitor = (0, [])
+dijkstra :: Grid Int -> Node -> Node -> (Node -> Int -> b) -> (Maybe Int, [b])
+dijkstra grid@Grid { items } start end visitor = go (M.singleton start 0) S.empty
+  where
+    visit x = let item = fromJust $ M.lookup x items in visitor x item
+    next ds seen = (M.lookupMin $ M.withoutKeys ds seen)
+    go ds seen = case next ds seen of
+        Nothing -> (Nothing, [])
+        Just (u, du)
+          | u == end -> (M.lookup u ds, [visit u])
+          | otherwise ->
+             let v = visit u
+                 seen' = S.insert u seen
+                 ds' = foldl (relax u du) ds (neighbours grid u)
+                 (d', vs) = go ds' seen'
+             in (d', v : vs)
+
+    relax :: Node -> Int -> M.Map Node Int -> Node -> M.Map Node Int
+    relax u du ds v = let d = distance grid u v in case M.lookup v ds of
+                         Just dv | dv < du + d -> ds
+                         _ -> M.insert v (du + d) ds
+