Last active
December 23, 2021 08:02
-
-
Save chowells79/1127f8bfe82f964cb6807cb869c055a9 to your computer and use it in GitHub Desktop.
Dijkstra's algorithm using psqueues package
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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