72 lines
2.5 KiB
Haskell
72 lines
2.5 KiB
Haskell
module Graph
|
|
( Graph (..),
|
|
Distance (..),
|
|
findShortestPath,
|
|
)
|
|
where
|
|
|
|
import qualified Data.HashMap.Strict as M
|
|
import Data.Hashable (Hashable)
|
|
import Data.Maybe (fromJust)
|
|
import qualified Data.PSQueue as PQ
|
|
|
|
data Distance a = Dist a | Infinity deriving (Eq)
|
|
|
|
instance (Ord a) => Ord (Distance a) where
|
|
Infinity <= Infinity = True
|
|
Infinity <= Dist _ = False
|
|
Dist _ <= Infinity = True
|
|
Dist x <= Dist y = x <= y
|
|
|
|
instance (Show a) => Show (Distance a) where
|
|
show Infinity = "Infinity"
|
|
show (Dist x) = show x
|
|
|
|
newtype Graph a b = Graph {edges :: M.HashMap a [(a, Distance b)]} deriving (Show)
|
|
|
|
addDistance :: (Num a) => Distance a -> Distance a -> Distance a
|
|
addDistance (Dist x) (Dist y) = Dist (x + y)
|
|
addDistance _ _ = Infinity
|
|
|
|
data DijkstraState a b = DijkstraState
|
|
{ unvisited :: PQ.PSQ a (Distance b),
|
|
distances :: M.HashMap a (Distance b)
|
|
}
|
|
|
|
updateDistances :: (Hashable a, Num b) => M.HashMap a (Distance b) -> [(a, Distance b)] -> Distance b -> M.HashMap a (Distance b)
|
|
updateDistances dists [] _ = dists
|
|
updateDistances dists (n : nodes) startD =
|
|
let newD = addDistance startD (snd n)
|
|
in updateDistances (M.adjust (const newD) (fst n) dists) nodes startD
|
|
|
|
visit :: (Ord a, Num b, Ord b) => PQ.PSQ a (Distance b) -> a -> [(a, Distance b)] -> Distance b -> PQ.PSQ a (Distance b)
|
|
visit us node [] _ = PQ.delete node us
|
|
visit us node (e : es) startD =
|
|
let newD = addDistance startD (snd e)
|
|
in visit (PQ.adjust (const newD) (fst e) us) node es startD
|
|
|
|
visitNode :: (Hashable a, Ord a, Num b, Ord b) => DijkstraState a b -> Graph a b -> a -> Distance b -> DijkstraState a b
|
|
visitNode state graph node d =
|
|
let es = edges graph M.! node
|
|
ds = updateDistances (distances state) es d
|
|
us = visit (unvisited state) node es d
|
|
in state {unvisited = us, distances = ds}
|
|
|
|
findShortestPath :: (Hashable a, Ord a, Ord b, Num b) => Graph a b -> a -> a -> Distance b
|
|
findShortestPath graph start end =
|
|
let nodesDist = (start PQ.:-> Dist 0) : [k PQ.:-> Infinity | k <- M.keys $ edges graph, k /= start]
|
|
dists = (start, Dist 0) : [(k, Infinity) | k <- M.keys $ edges graph, k /= start]
|
|
initialState = DijkstraState {unvisited = PQ.fromList nodesDist, distances = M.fromList dists}
|
|
in dijkstra initialState
|
|
where
|
|
dijkstra s =
|
|
let nd = fromJust $ PQ.findMin (unvisited s)
|
|
n = PQ.key nd
|
|
d = PQ.prio nd
|
|
in if n == end
|
|
then d
|
|
else
|
|
if d == Infinity
|
|
then Infinity
|
|
else dijkstra $ visitNode s graph n d
|