{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ScopedTypeVariables #-}
module ToySolver.Graph.ShortestPath
(
Graph
, Edge
, OutEdge
, InEdge
, Fold (..)
, monoid'
, monoid
, unit
, pair
, path
, firstOutEdge
, lastInEdge
, cost
, Path (..)
, pathFrom
, pathTo
, pathCost
, pathEmpty
, pathAppend
, pathEdges
, pathEdgesBackward
, pathEdgesSeq
, pathVertexes
, pathVertexesBackward
, pathVertexesSeq
, pathFold
, pathMin
, bellmanFord
, dijkstra
, floydWarshall
, bellmanFordDetectNegativeCycle
) where
import Control.Monad
import Control.Monad.ST
import Control.Monad.Trans
import Control.Monad.Trans.Except
import Data.Hashable
import qualified Data.HashTable.Class as H
import qualified Data.HashTable.ST.Cuckoo as C
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import qualified Data.HashSet as HashSet
import qualified Data.Heap as Heap
import Data.List (foldl')
import Data.Monoid
import Data.Ord
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import Data.STRef
type Graph vertex cost label = HashMap vertex [OutEdge vertex cost label]
type Edge vertex cost label = (vertex, vertex, cost, label)
type OutEdge vertex cost label = (vertex, cost, label)
type InEdge vertex cost label = (vertex, cost, label)
data Path vertex cost label
= Empty vertex
| Singleton (Edge vertex cost label)
| Append (Path vertex cost label) (Path vertex cost label) !cost
deriving (Eq, Show)
pathFrom :: Path vertex cost label -> vertex
pathFrom (Empty v) = v
pathFrom (Singleton (from,_,_,_)) = from
pathFrom (Append p1 _ _) = pathFrom p1
pathTo :: Path vertex cost label -> vertex
pathTo (Empty v) = v
pathTo (Singleton (_,to,_,_)) = to
pathTo (Append _ p2 _) = pathTo p2
pathCost :: Num cost => Path vertex cost label -> cost
pathCost (Empty _) = 0
pathCost (Singleton (_,_,c,_)) = c
pathCost (Append _ _ c) = c
pathEmpty :: vertex -> Path vertex cost label
pathEmpty = Empty
pathAppend :: (Eq vertex, Num cost) => Path vertex cost label -> Path vertex cost label -> Path vertex cost label
pathAppend p1 p2
| pathTo p1 /= pathFrom p2 = error "ToySolver.Graph.ShortestPath.pathAppend: pathTo/pathFrom mismatch"
| otherwise =
case (p1, p2) of
(Empty _, _) -> p2
(_, Empty _) -> p1
_ -> Append p1 p2 (pathCost p1 + pathCost p2)
pathEdges :: Path vertex cost label -> [Edge vertex cost label]
pathEdges p = f p []
where
f (Empty _) xs = xs
f (Singleton e) xs = e : xs
f (Append p1 p2 _) xs = f p1 (f p2 xs)
pathEdgesBackward :: Path vertex cost label -> [Edge vertex cost label]
pathEdgesBackward p = f p []
where
f (Empty _) xs = xs
f (Singleton e) xs = e : xs
f (Append p1 p2 _) xs = f p2 (f p1 xs)
pathEdgesSeq :: Path vertex cost label -> Seq (Edge vertex cost label)
pathEdgesSeq (Empty _) = Seq.empty
pathEdgesSeq (Singleton e) = Seq.singleton e
pathEdgesSeq (Append p1 p2 _) = pathEdgesSeq p1 <> pathEdgesSeq p2
pathVertexes :: Path vertex cost label -> [vertex]
pathVertexes p = pathFrom p : f p []
where
f (Empty _) xs = xs
f (Singleton (_,v2,_,_)) xs = v2 : xs
f (Append p1 p2 _) xs = f p1 (f p2 xs)
pathVertexesBackward :: Path vertex cost label -> [vertex]
pathVertexesBackward p = pathTo p : f p []
where
f (Empty _) xs = xs
f (Singleton (v1,_,_,_)) xs = v1 : xs
f (Append p1 p2 _) xs = f p2 (f p1 xs)
pathVertexesSeq :: Path vertex cost label -> Seq vertex
pathVertexesSeq p = f True p
where
f True (Empty v) = Seq.singleton v
f False (Empty _) = mempty
f True (Singleton (v1,v2,_,_)) = Seq.fromList [v1, v2]
f False (Singleton (v1,_,_,_)) = Seq.singleton v1
f b (Append p1 p2 _) = f False p1 <> f b p2
pathMin :: (Num cost, Ord cost) => Path vertex cost label -> Path vertex cost label -> Path vertex cost label
pathMin p1 p2
| pathCost p1 <= pathCost p2 = p1
| otherwise = p2
pathFold :: Fold vertex cost label a -> Path vertex cost label -> a
pathFold (Fold fV fE fC fD) p = fD (h p)
where
h (Empty v) = fV v
h (Singleton e) = fE e
h (Append p1 p2 _) = fC (h p1) (h p2)
data Pair a b = Pair !a !b
data Fold vertex cost label r
= forall a. Fold (vertex -> a) (Edge vertex cost label -> a) (a -> a -> a) (a -> r)
instance Functor (Fold vertex cost label) where
{-# INLINE fmap #-}
fmap f (Fold fV fE fC fD) = Fold fV fE fC (f . fD)
instance Applicative (Fold vertex cost label) where
{-# INLINE pure #-}
pure a = Fold (\_ -> ()) (\_ -> ()) (\_ _ -> ()) (const a)
{-# INLINE (<*>) #-}
Fold fV1 fE1 fC1 fD1 <*> Fold fV2 fE2 fC2 fD2 =
Fold (\v -> Pair (fV1 v) (fV2 v))
(\e -> Pair (fE1 e) (fE2 e))
(\(Pair a1 b1) (Pair a2 b2) -> Pair (fC1 a1 a2) (fC2 b1 b2))
(\(Pair a b) -> fD1 a (fD2 b))
monoid' :: Monoid m => (Edge vertex cost label -> m) -> Fold vertex cost label m
monoid' f = Fold (\_ -> mempty) f mappend id
monoid :: Monoid m => Fold vertex cost m m
monoid = monoid' (\(_,_,_,m) -> m)
unit :: Fold vertex cost label ()
unit = monoid' (\_ -> ())
pair :: Fold vertex cost label a -> Fold vertex cost label b -> Fold vertex cost label (a,b)
pair (Fold fV1 fE1 fC1 fD1) (Fold fV2 fE2 fC2 fD2) =
Fold (\v -> Pair (fV1 v) (fV2 v))
(\e -> Pair (fE1 e) (fE2 e))
(\(Pair a1 b1) (Pair a2 b2) -> Pair (fC1 a1 a2) (fC2 b1 b2))
(\(Pair a b) -> (fD1 a, fD2 b))
path :: (Eq vertex, Num cost) => Fold vertex cost label (Path vertex cost label)
path = Fold pathEmpty Singleton pathAppend id
cost :: Num cost => Fold vertex cost label cost
cost = Fold (\_ -> 0) (\(_,_,c,_) -> c) (+) id
firstOutEdge :: Fold vertex cost label (First (OutEdge vertex cost label))
firstOutEdge = monoid' (\(_,v,c,l) -> First (Just (v,c,l)))
lastInEdge :: Fold vertex cost label (Last (InEdge vertex cost label))
lastInEdge = monoid' (\(v,_,c,l) -> Last (Just (v,c,l)))
bellmanFord
:: (Eq vertex, Hashable vertex, Real cost)
=> Fold vertex cost label a
-> Graph vertex cost label
-> [vertex]
-> HashMap vertex (cost, a)
bellmanFord (Fold fV fE fC fD) g ss = runST $ do
let n = HashMap.size g
d <- C.newSized n
forM_ ss $ \s -> H.insert d s (Pair 0 (fV s))
updatedRef <- newSTRef (HashSet.fromList ss)
_ <- runExceptT $ replicateM_ n $ do
us <- lift $ readSTRef updatedRef
when (HashSet.null us) $ throwE ()
lift $ do
writeSTRef updatedRef HashSet.empty
forM_ (HashSet.toList us) $ \u -> do
Just (Pair du a) <- H.lookup d u
forM_ (HashMap.lookupDefault [] u g) $ \(v, c, l) -> do
m <- H.lookup d v
case m of
Just (Pair c0 _) | c0 <= du + c -> return ()
_ -> do
H.insert d v (Pair (du + c) (a `fC` fE (u,v,c,l)))
modifySTRef' updatedRef (HashSet.insert v)
liftM (fmap (\(Pair c x) -> (c, fD x))) $ freezeHashTable d
freezeHashTable :: (H.HashTable h, Hashable k, Eq k) => h s k v -> ST s (HashMap k v)
freezeHashTable h = H.foldM (\m (k,v) -> return $! HashMap.insert k v m) HashMap.empty h
bellmanFordDetectNegativeCycle
:: forall vertex cost label a. (Eq vertex, Hashable vertex, Real cost)
=> Fold vertex cost label a
-> Graph vertex cost label
-> HashMap vertex (cost, Last (InEdge vertex cost label))
-> Maybe a
bellmanFordDetectNegativeCycle (Fold fV fE fC fD) g d = either (Just . fD) (const Nothing) $ do
forM_ (HashMap.toList d) $ \(u,(du,_)) ->
forM_ (HashMap.lookupDefault [] u g) $ \(v, c, l) -> do
let (dv,_) = d HashMap.! v
when (du + c < dv) $ do
let d' = HashMap.insert v (du + c, Last (Just (u, c, l))) d
parent u = do
case HashMap.lookup u d' of
Just (_, Last (Just (v,_,_))) -> v
_ -> undefined
u0 = go (parent (parent v)) (parent v)
where
go hare tortoise
| hare == tortoise = hare
| otherwise = go (parent (parent hare)) (parent tortoise)
let go u result = do
let Just (_, Last (Just (v,c,l))) = HashMap.lookup u d'
if v == u0 then
fC (fE (v,u,c,l)) result
else
go v (fC (fE (v,u,c,l)) result)
Left $ go u0 (fV u0)
dijkstra
:: forall vertex cost label a. (Eq vertex, Hashable vertex, Real cost)
=> Fold vertex cost label a
-> Graph vertex cost label
-> [vertex]
-> HashMap vertex (cost, a)
dijkstra (Fold fV fE fC (fD :: x -> a)) g ss =
fmap (\(Pair c x) -> (c, fD x)) $
loop (Heap.fromList [Heap.Entry 0 (Pair s (fV s)) | s <- ss]) HashMap.empty
where
loop
:: Heap.Heap (Heap.Entry cost (Pair vertex x))
-> HashMap vertex (Pair cost x)
-> HashMap vertex (Pair cost x)
loop q visited =
case Heap.viewMin q of
Nothing -> visited
Just (Heap.Entry c (Pair v a), q')
| v `HashMap.member` visited -> loop q' visited
| otherwise ->
let q2 = Heap.fromList
[ Heap.Entry (c+c') (Pair ch (a `fC` fE (v,ch,c',l)))
| (ch,c',l) <- HashMap.lookupDefault [] v g
, not (ch `HashMap.member` visited)
]
in loop (Heap.union q' q2) (HashMap.insert v (Pair c a) visited)
floydWarshall
:: forall vertex cost label a. (Eq vertex, Hashable vertex, Real cost)
=> Fold vertex cost label a
-> Graph vertex cost label
-> HashMap vertex (HashMap vertex (cost, a))
floydWarshall (Fold fV fE fC (fD :: x -> a)) g =
fmap (fmap (\(Pair c x) -> (c, fD x))) $
HashMap.unionWith (HashMap.unionWith minP) (foldl' f tbl0 vs) paths0
where
vs = HashMap.keys g
paths0 :: HashMap vertex (HashMap vertex (Pair cost x))
paths0 = HashMap.mapWithKey (\v _ -> HashMap.singleton v (Pair 0 (fV v))) g
tbl0 :: HashMap vertex (HashMap vertex (Pair cost x))
tbl0 = HashMap.mapWithKey (\v es -> HashMap.fromListWith minP [(u, (Pair c (fE (v,u,c,l)))) | (u,c,l) <- es]) g
minP :: Pair cost x -> Pair cost x -> Pair cost x
minP = minBy (comparing (\(Pair c _) -> c))
f :: HashMap vertex (HashMap vertex (Pair cost x))
-> vertex
-> HashMap vertex (HashMap vertex (Pair cost x))
f tbl vk =
case HashMap.lookup vk tbl of
Nothing -> tbl
Just hk -> HashMap.map h tbl
where
h :: HashMap vertex (Pair cost x) -> HashMap vertex (Pair cost x)
h m =
case HashMap.lookup vk m of
Nothing -> m
Just (Pair c1 x1) -> HashMap.unionWith minP m (HashMap.map (\(Pair c2 x2) -> (Pair (c1+c2) (fC x1 x2))) hk)
minBy :: (a -> a -> Ordering) -> a -> a -> a
minBy f a b =
case f a b of
LT -> a
EQ -> a
GT -> b