Skip to content

Instantly share code, notes, and snippets.

@chowells79
Last active December 23, 2021 08:02
Show Gist options
  • Save chowells79/1127f8bfe82f964cb6807cb869c055a9 to your computer and use it in GitHub Desktop.
Save chowells79/1127f8bfe82f964cb6807cb869c055a9 to your computer and use it in GitHub Desktop.
Dijkstra's algorithm using psqueues package
import Data.List (foldl')
-- unordered-containers
import Data.Hashable (Hashable)
import qualified Data.HashMap.Strict as HM
-- psqueues
import qualified Data.HashPSQ as P
-- Streams all nodes as they are found in order of non-decreasing cost
-- In addition to the located node, also returns the path used
-- The first argument is provided the current cost and node and should
-- return the nodes reachable from there along with the cumulative cost
-- to reach them from this location, given the current cost.
-- For simpler versions, see dijkstraMonoid or dijkstraSum
dijkstra
:: (Ord cost, Ord node, Hashable node)
=> (cost -> node -> [(cost, node)]) -> cost -> node
-> [(node, (cost, [node]))]
dijkstra neighbors startC startN =
go (HM.singleton startN [startN]) P.empty startC startN
where
go paths reachable cost node = case P.minView reachable' of
Nothing -> []
Just (next, total', link, reachable'') ->
report : go paths' reachable'' total' next
where
report = (node, (cost, reverse $ paths HM.! node))
paths' = HM.insert next (next : paths HM.! link) paths
where
near = filter unvisited $ neighbors cost node
unvisited (_, x) = not $ HM.member x paths
reachable' = foldl' (\r (c, n) -> lower n c r) reachable near
lower k p = snd . P.alter minP k
where
minP Nothing = ((), Just (p, node))
minP (Just e@(p', _)) =
((), Just $ if p < p' then (p, node) else e)
-- uses Monoidal composition to combine local costs with the
-- cumulative cost, setting the initial cost to mempty
dijkstraMonoid
:: (Monoid cost, Ord cost, Ord node, Hashable node)
=> (node -> [(cost, node)]) -> node -> [(node, (cost, [node]))]
dijkstraMonoid f n =
dijkstra (\c d -> map (\(x, y) -> (c <> x, y)) (f d)) mempty n
-- Uses addition to combine local costs with cumulative cost, setting
-- the initial cost to 0
dijkstraSum
:: (Num cost, Ord cost, Ord node, Hashable node)
=> (node -> [(cost, node)]) -> node -> [(node, (cost, [node]))]
dijkstraSum f n = dijkstra (\c d -> map (\(x, y) -> (c + x, y)) (f d)) 0 n
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment