-- SPDX-License-Identifier: MPL-2.0

{- | Checked graph paths. Unreachable distances are 'Nothing'. A reachable
negative cycle is a domain error; all-pairs methods reject any negative cycle.
Nonfinite arithmetic is rejected. Dijkstra-family methods require nonnegative
weights. Returned vertex paths use lexicographic tie-breaking where applicable.
-}
module HMath.Graph.Paths (
    dijkstra,
    multiSourceDijkstra,
    bellmanFord,
    floydWarshall,
    johnson,
    aStar,
    bidirectionalDijkstra,
    kShortestPaths,
    longestDAGPaths,
) where

import Control.Monad (foldM)
import qualified Data.Foldable as Foldable (foldl')
import Data.List (sortOn)
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import HMath.Error (MathError (..))
import HMath.Graph.Core

type Distances = M.Map Int Double

finish :: Graph -> Distances -> Either MathError [Maybe Double]
finish :: Graph -> Distances -> Either MathError [Maybe Double]
finish Graph
g Distances
values
    | (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (Double -> Bool) -> Double -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Bool
finite) (Distances -> [Double]
forall k a. Map k a -> [a]
M.elems Distances
values) = MathError -> Either MathError [Maybe Double]
forall a b. a -> Either a b
Left MathError
NonFiniteResult
    | Bool
otherwise = [Maybe Double] -> Either MathError [Maybe Double]
forall a b. b -> Either a b
Right [Int -> Distances -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
v Distances
values | Int
v <- [Int
0 .. Graph -> Int
vertexCount Graph
g Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
nonnegative :: Graph -> Either MathError ()
nonnegative :: Graph -> Either MathError ()
nonnegative Graph
g
    | (Edge -> Bool) -> [Edge] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0) (Double -> Bool) -> (Edge -> Double) -> Edge -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Edge -> Double
edgeWeight) (Graph -> [Edge]
edges Graph
g) = MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"nonnegative edge weights required")
    | Bool
otherwise = () -> Either MathError ()
forall a b. b -> Either a b
Right ()

checkedSum :: Double -> Double -> Either MathError Double
checkedSum :: Double -> Double -> Either MathError Double
checkedSum Double
a Double
b = let c :: Double
c = Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b in if Double -> Bool
finite Double
c then Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
c else MathError -> Either MathError Double
forall a b. a -> Either a b
Left MathError
NonFiniteResult

{- | Given graph then source list, return distance to the nearest supplied source, each
seeded at zero. Invalid sources, negative weights or non-finite arithmetic fail. Empty
sources give all Nothing.
-}
multiSourceDijkstra :: Graph -> [Int] -> Either MathError [Maybe Double]
multiSourceDijkstra :: Graph -> [Int] -> Either MathError [Maybe Double]
multiSourceDijkstra Graph
g [Int]
sources = do
    (Int -> Either MathError ()) -> [Int] -> Either MathError ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Graph -> Int -> Either MathError ()
checkVertex Graph
g) [Int]
sources
    Graph -> Either MathError ()
nonnegative Graph
g
    distances <- Set Int -> Distances -> Either MathError Distances
go Set Int
forall a. Set a
S.empty ([(Int, Double)] -> Distances
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
v, Double
0) | Int
v <- [Int]
sources])
    finish g distances
  where
    go :: Set Int -> Distances -> Either MathError Distances
go Set Int
visited Distances
dist = case [(Double
d, Int
v) | (Int
v, Double
d) <- Distances -> [(Int, Double)]
forall k a. Map k a -> [(k, a)]
M.toList Distances
dist, Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
S.notMember Int
v Set Int
visited] of
        [] -> Distances -> Either MathError Distances
forall a b. b -> Either a b
Right Distances
dist
        [(Double, Int)]
options -> do
            let (Double
_, Int
v) = [(Double, Int)] -> (Double, Int)
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [(Double, Int)]
options
            next <-
                (Distances -> (Int, Double, Int) -> Either MathError Distances)
-> Distances -> [(Int, Double, Int)] -> Either MathError Distances
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
                    (\Distances
m (Int
u, Double
w, Int
_) -> do candidate <- Double -> Double -> Either MathError Double
checkedSum (Distances
dist Distances -> Int -> Double
forall k a. Ord k => Map k a -> k -> a
M.! Int
v) Double
w; pure (M.insertWith min u candidate m))
                    Distances
dist
                    (Graph -> Int -> [(Int, Double, Int)]
neighbors Graph
g Int
v)
            go (S.insert v visited) next

{- | Given graph then source, return minimum distances in vertex order, with Nothing for
unreachable vertices. Require valid source and nonnegative weights; invalid input or
non-finite arithmetic fails.
-}
dijkstra :: Graph -> Int -> Either MathError [Maybe Double]
dijkstra :: Graph -> Int -> Either MathError [Maybe Double]
dijkstra Graph
g Int
source = Graph -> [Int] -> Either MathError [Maybe Double]
multiSourceDijkstra Graph
g [Int
source]

relax :: [Edge] -> Distances -> Either MathError Distances
relax :: [Edge] -> Distances -> Either MathError Distances
relax [Edge]
es Distances
initial = (Distances -> Edge -> Either MathError Distances)
-> Distances -> [Edge] -> Either MathError Distances
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Distances -> Edge -> Either MathError Distances
step Distances
initial [Edge]
es
  where
    step :: Distances -> Edge -> Either MathError Distances
step Distances
m Edge
e = case Int -> Distances -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Edge -> Int
edgeFrom Edge
e) Distances
m of
        Maybe Double
Nothing -> Distances -> Either MathError Distances
forall a b. b -> Either a b
Right Distances
m
        Just Double
d -> do candidate <- Double -> Double -> Either MathError Double
checkedSum Double
d (Edge -> Double
edgeWeight Edge
e); pure (M.insertWith min (edgeTo e) candidate m)
repeatRelax :: Int -> [Edge] -> Distances -> Either MathError Distances
repeatRelax :: Int -> [Edge] -> Distances -> Either MathError Distances
repeatRelax Int
count [Edge]
es Distances
initial = (Distances -> Int -> Either MathError Distances)
-> Distances -> [Int] -> Either MathError Distances
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Distances
m Int
_ -> [Edge] -> Distances -> Either MathError Distances
relax [Edge]
es Distances
m) Distances
initial [Int
1 .. Int
count]

{- | Given graph then source, return minimum distances including negative weights. Invalid
source, reachable negative cycle or non-finite arithmetic returns Left MathError;
unreachable vertices are Nothing.
-}
bellmanFord :: Graph -> Int -> Either MathError [Maybe Double]
bellmanFord :: Graph -> Int -> Either MathError [Maybe Double]
bellmanFord Graph
g Int
source = do
    Graph -> Int -> Either MathError ()
checkVertex Graph
g Int
source
    dist <- Int -> [Edge] -> Distances -> Either MathError Distances
repeatRelax (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Graph -> Int
vertexCount Graph
g Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (Graph -> [Edge]
arcs Graph
g) (Int -> Double -> Distances
forall k a. k -> a -> Map k a
M.singleton Int
source Double
0)
    next <- relax (arcs g) dist
    if next /= dist then Left (DomainError "reachable negative cycle") else finish g dist

{- | Return all-pairs minimum-distance rows, with Nothing for unreachable pairs. Negative
edges are allowed; any negative cycle or non-finite arithmetic fails.
-}
floydWarshall :: Graph -> Either MathError [[Maybe Double]]
floydWarshall :: Graph -> Either MathError [[Maybe Double]]
floydWarshall Graph
g = do
    let n :: Int
n = Graph -> Int
vertexCount Graph
g
        initial :: Map (Int, Int) Double
initial =
            (Double -> Double -> Double)
-> [((Int, Int), Double)] -> Map (Int, Int) Double
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith
                Double -> Double -> Double
forall a. Ord a => a -> a -> a
min
                ([((Int
v, Int
v), Double
0) | Int
v <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] [((Int, Int), Double)]
-> [((Int, Int), Double)] -> [((Int, Int), Double)]
forall a. [a] -> [a] -> [a]
++ [((Edge -> Int
edgeFrom Edge
e, Edge -> Int
edgeTo Edge
e), Edge -> Double
edgeWeight Edge
e) | Edge
e <- Graph -> [Edge]
arcs Graph
g])
        update :: Map (Int, Int) Double
-> Int -> Either MathError (Map (Int, Int) Double)
update Map (Int, Int) Double
dist Int
k =
            (Map (Int, Int) Double
 -> (Int, Int) -> Either MathError (Map (Int, Int) Double))
-> Map (Int, Int) Double
-> [(Int, Int)]
-> Either MathError (Map (Int, Int) Double)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
                ( \Map (Int, Int) Double
m (Int
i, Int
j) -> case ((Int, Int) -> Map (Int, Int) Double -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Int
i, Int
k) Map (Int, Int) Double
dist, (Int, Int) -> Map (Int, Int) Double -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Int
k, Int
j) Map (Int, Int) Double
dist) of
                    (Just Double
a, Just Double
b) -> do candidate <- Double -> Double -> Either MathError Double
checkedSum Double
a Double
b; pure (M.insertWith min (i, j) candidate m)
                    (Maybe Double, Maybe Double)
_ -> Map (Int, Int) Double -> Either MathError (Map (Int, Int) Double)
forall a b. b -> Either a b
Right Map (Int, Int) Double
m
                )
                Map (Int, Int) Double
dist
                [(Int
i, Int
j) | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    out <- (Map (Int, Int) Double
 -> Int -> Either MathError (Map (Int, Int) Double))
-> Map (Int, Int) Double
-> [Int]
-> Either MathError (Map (Int, Int) Double)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Map (Int, Int) Double
-> Int -> Either MathError (Map (Int, Int) Double)
update Map (Int, Int) Double
initial [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    if any (\Int
i -> Double -> (Int, Int) -> Map (Int, Int) Double -> Double
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Double
0 (Int
i, Int
i) Map (Int, Int) Double
out Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0) [0 .. n - 1]
        then Left (DomainError "negative cycle")
        else pure [[M.lookup (i, j) out | j <- [0 .. n - 1]] | i <- [0 .. n - 1]]

{- | Johnson reweighting: Bellman--Ford potentials from an implicit zero-cost
super-source, followed by independent Dijkstra searches.
-}
johnson :: Graph -> Either MathError [[Maybe Double]]
johnson :: Graph -> Either MathError [[Maybe Double]]
johnson Graph
g = do
    let n :: Int
n = Graph -> Int
vertexCount Graph
g
    potentials <- Int -> [Edge] -> Distances -> Either MathError Distances
repeatRelax Int
n (Graph -> [Edge]
arcs Graph
g) ([(Int, Double)] -> Distances
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
v, Double
0) | Int
v <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
    next <- relax (arcs g) potentials
    if next /= potentials
        then Left (DomainError "negative cycle")
        else do
            triples <- mapM (reweight potentials) (arcs g)
            weighted <- graph True n triples
            rows <- mapM (dijkstra weighted) [0 .. n - 1]
            mapM
                ( \Int
i ->
                    (Int -> Either MathError (Maybe Double))
-> [Int] -> Either MathError [Maybe Double]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
                        ( \Int
j -> case [[Maybe Double]]
rows [[Maybe Double]] -> Int -> [Maybe Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [Maybe Double] -> Int -> Maybe Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j of
                            Maybe Double
Nothing -> Maybe Double -> Either MathError (Maybe Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Double
forall a. Maybe a
Nothing
                            Just Double
distance -> do a <- Double -> Double -> Either MathError Double
checkedSum Double
distance (-(Distances
potentials Distances -> Int -> Double
forall k a. Ord k => Map k a -> k -> a
M.! Int
i)); Just <$> checkedSum a (potentials M.! j)
                        )
                        [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                )
                [0 .. n - 1]
  where
    reweight :: Distances -> Edge -> Either MathError (Int, Int, Double)
reweight Distances
potentials Edge
e = do
        a <- Double -> Double -> Either MathError Double
checkedSum (Edge -> Double
edgeWeight Edge
e) (Distances
potentials Distances -> Int -> Double
forall k a. Ord k => Map k a -> k -> a
M.! Edge -> Int
edgeFrom Edge
e)
        x <- checkedSum a (-(potentials M.! edgeTo e))
        let scale = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1 ([Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
abs [Edge -> Double
edgeWeight Edge
e, Distances
potentials Distances -> Int -> Double
forall k a. Ord k => Map k a -> k -> a
M.! Edge -> Int
edgeFrom Edge
e, Distances
potentials Distances -> Int -> Double
forall k a. Ord k => Map k a -> k -> a
M.! Edge -> Int
edgeTo Edge
e]))
        -- Clamp only cancellation within the stated relative rounding scale.
        if x < -1e-12 * scale
            then Left (DomainError "unstable Johnson reweighting")
            else pure (edgeFrom e, edgeTo e, max 0 x)

{- | A* requires a finite consistent heuristic, h(target)=0, checked on all
arcs. This guarantees optimality without trusting an arbitrary heuristic.
-}
aStar :: Graph -> (Int -> Double) -> Int -> Int -> Either MathError (Maybe (Double, [Int]))
aStar :: Graph
-> (Int -> Double)
-> Int
-> Int
-> Either MathError (Maybe (Double, [Int]))
aStar Graph
g Int -> Double
heuristic Int
source Int
target = do
    Graph -> Int -> Either MathError ()
checkVertex Graph
g Int
source
    Graph -> Int -> Either MathError ()
checkVertex Graph
g Int
target
    Graph -> Either MathError ()
nonnegative Graph
g
    (Edge -> Either MathError ()) -> [Edge] -> Either MathError ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Edge
e -> Double -> Double -> Either MathError Double
checkedSum (Edge -> Double
edgeWeight Edge
e) (Int -> Double
heuristic (Edge -> Int
edgeTo Edge
e)) Either MathError Double
-> Either MathError () -> Either MathError ()
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> Either MathError ()
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (Graph -> [Edge]
arcs Graph
g)
    let hs :: [Double]
hs = [Int -> Double
heuristic Int
v | Int
v <- [Int
0 .. Graph -> Int
vertexCount Graph
g Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    if (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (Double -> Bool) -> Double -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Bool
finite) [Double]
hs
        Bool -> Bool -> Bool
|| Int -> Double
heuristic Int
target Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Double
0
        Bool -> Bool -> Bool
|| (Edge -> Bool) -> [Edge] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
            ( \Edge
e ->
                Int -> Double
heuristic (Edge -> Int
edgeFrom Edge
e)
                    Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Edge -> Double
edgeWeight Edge
e Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
heuristic (Edge -> Int
edgeTo Edge
e) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1e-12 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1 (Double -> Double
forall a. Num a => a -> a
abs (Int -> Double
heuristic (Edge -> Int
edgeFrom Edge
e)))
            )
            (Graph -> [Edge]
arcs Graph
g)
        then MathError -> Either MathError (Maybe (Double, [Int]))
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"A* heuristic must be finite and consistent")
        else Set Int
-> Map Int (Double, [Int])
-> Either MathError (Maybe (Double, [Int]))
go Set Int
forall a. Set a
S.empty (Int -> (Double, [Int]) -> Map Int (Double, [Int])
forall k a. k -> a -> Map k a
M.singleton Int
source (Double
0, [Int
source]))
  where
    go :: Set Int
-> Map Int (Double, [Int])
-> Either MathError (Maybe (Double, [Int]))
go Set Int
closed Map Int (Double, [Int])
known = do
        candidates <-
            ((Int, (Double, [Int]))
 -> Either MathError (Double, Double, [Int], Int))
-> [(Int, (Double, [Int]))]
-> Either MathError [(Double, Double, [Int], Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
                (\(Int
v, (Double
d, [Int]
path)) -> do priority <- Double -> Double -> Either MathError Double
checkedSum Double
d (Int -> Double
heuristic Int
v); pure (priority, d, path, v))
                [(Int
v, (Double, [Int])
p) | (Int
v, (Double, [Int])
p) <- Map Int (Double, [Int]) -> [(Int, (Double, [Int]))]
forall k a. Map k a -> [(k, a)]
M.toList Map Int (Double, [Int])
known, Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
S.notMember Int
v Set Int
closed]
        case candidates of
            [] -> Maybe (Double, [Int]) -> Either MathError (Maybe (Double, [Int]))
forall a b. b -> Either a b
Right Maybe (Double, [Int])
forall a. Maybe a
Nothing
            [(Double, Double, [Int], Int)]
_ -> do
                let (Double
_, Double
d, [Int]
path, Int
v) = [(Double, Double, [Int], Int)] -> (Double, Double, [Int], Int)
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [(Double, Double, [Int], Int)]
candidates
                if Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
target
                    then Maybe (Double, [Int]) -> Either MathError (Maybe (Double, [Int]))
forall a b. b -> Either a b
Right ((Double, [Int]) -> Maybe (Double, [Int])
forall a. a -> Maybe a
Just (Double
d, [Int]
path))
                    else do
                        next <-
                            (Map Int (Double, [Int])
 -> (Int, Double, Int)
 -> Either MathError (Map Int (Double, [Int])))
-> Map Int (Double, [Int])
-> [(Int, Double, Int)]
-> Either MathError (Map Int (Double, [Int]))
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
                                (\Map Int (Double, [Int])
m (Int
u, Double
w, Int
_) -> do candidate <- Double -> Double -> Either MathError Double
checkedSum Double
d Double
w; pure (M.insertWith min u (candidate, path ++ [u]) m))
                                Map Int (Double, [Int])
known
                                (Graph -> Int -> [(Int, Double, Int)]
neighbors Graph
g Int
v)
                        go (S.insert v closed) next

{- | Two settled-frontier Dijkstra searches, alternating the frontier with the
smaller minimum key. The best meeting bound stops the search safely.
-}
bidirectionalDijkstra :: Graph -> Int -> Int -> Either MathError (Maybe Double)
bidirectionalDijkstra :: Graph -> Int -> Int -> Either MathError (Maybe Double)
bidirectionalDijkstra Graph
g Int
source Int
target = do
    Graph -> Int -> Either MathError ()
checkVertex Graph
g Int
source
    Graph -> Int -> Either MathError ()
checkVertex Graph
g Int
target
    Graph -> Either MathError ()
nonnegative Graph
g
    Distances
-> Distances
-> Set Int
-> Set Int
-> Maybe Double
-> Either MathError (Maybe Double)
go (Int -> Double -> Distances
forall k a. k -> a -> Map k a
M.singleton Int
source Double
0) (Int -> Double -> Distances
forall k a. k -> a -> Map k a
M.singleton Int
target Double
0) Set Int
forall a. Set a
S.empty Set Int
forall a. Set a
S.empty Maybe Double
forall a. Maybe a
Nothing
  where
    reverseNeighbors :: Int -> [(Int, Double)]
reverseNeighbors Int
v = [(Edge -> Int
edgeFrom Edge
e, Edge -> Double
edgeWeight Edge
e) | Edge
e <- Graph -> [Edge]
arcs Graph
g, Edge -> Int
edgeTo Edge
e Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
v]
    minimumOpen :: Map b a -> Set b -> Maybe (a, b)
minimumOpen Map b a
dist Set b
closed = case [(a
d, b
v) | (b
v, a
d) <- Map b a -> [(b, a)]
forall k a. Map k a -> [(k, a)]
M.toList Map b a
dist, b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
S.notMember b
v Set b
closed] of
        [] -> Maybe (a, b)
forall a. Maybe a
Nothing
        [(a, b)]
xs -> (a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just ([(a, b)] -> (a, b)
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [(a, b)]
xs)
    go :: Distances
-> Distances
-> Set Int
-> Set Int
-> Maybe Double
-> Either MathError (Maybe Double)
go Distances
forward Distances
backward Set Int
sf Set Int
sb Maybe Double
best = case (Distances -> Set Int -> Maybe (Double, Int)
forall {b} {a}. (Ord b, Ord a) => Map b a -> Set b -> Maybe (a, b)
minimumOpen Distances
forward Set Int
sf, Distances -> Set Int -> Maybe (Double, Int)
forall {b} {a}. (Ord b, Ord a) => Map b a -> Set b -> Maybe (a, b)
minimumOpen Distances
backward Set Int
sb) of
        (Maybe (Double, Int)
Nothing, Maybe (Double, Int)
_) -> Maybe Double -> Either MathError (Maybe Double)
forall a b. b -> Either a b
Right Maybe Double
best
        (Maybe (Double, Int)
_, Maybe (Double, Int)
Nothing) -> Maybe Double -> Either MathError (Maybe Double)
forall a b. b -> Either a b
Right Maybe Double
best
        (Just (Double
a, Int
u), Just (Double
b, Int
v)) -> do
            lowerBound <- Double -> Double -> Either MathError Double
checkedSum Double
a Double
b
            if maybe False (<= lowerBound) best
                then Right best
                else
                    if a <= b
                        then expand True u a forward backward sf sb best
                        else expand False v b backward forward sb sf best
    expand :: Bool
-> Int
-> Double
-> Distances
-> Distances
-> Set Int
-> Set Int
-> Maybe Double
-> Either MathError (Maybe Double)
expand Bool
front Int
v Double
d Distances
own Distances
other Set Int
settled Set Int
opposite Maybe Double
best = do
        let links :: [(Int, Double)]
links = if Bool
front then [(Int
u, Double
w) | (Int
u, Double
w, Int
_) <- Graph -> Int -> [(Int, Double, Int)]
neighbors Graph
g Int
v] else Int -> [(Int, Double)]
reverseNeighbors Int
v
        next <-
            (Distances -> (Int, Double) -> Either MathError Distances)
-> Distances -> [(Int, Double)] -> Either MathError Distances
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Distances
m (Int
u, Double
w) -> do candidate <- Double -> Double -> Either MathError Double
checkedSum Double
d Double
w; pure (M.insertWith min u candidate m)) Distances
own [(Int, Double)]
links
        meetings <-
            mapM (uncurry checkedSum) [(x, y) | (u, x) <- M.toList next, Just y <- [M.lookup u other]]
        let updated = (Maybe Double -> Double -> Maybe Double)
-> Maybe Double -> [Double] -> Maybe Double
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\Maybe Double
old Double
x -> Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> (Double -> Double) -> Maybe Double -> Double
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Double
x (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
x) Maybe Double
old)) Maybe Double
best [Double]
meetings
        if front
            then go next other (S.insert v settled) opposite updated
            else go other next opposite (S.insert v settled) updated

{- | K cheapest simple vertex paths, best-first expansion, for nonnegative
weights. Parallel edges generate distinct edge paths; output includes edge
IDs to preserve this identity. Fuel counts expanded partial paths.
-}
kShortestPaths :: Int -> Int -> Graph -> Int -> Int -> Either MathError [(Double, [Int], [Int])]
kShortestPaths :: Int
-> Int
-> Graph
-> Int
-> Int
-> Either MathError [(Double, [Int], [Int])]
kShortestPaths Int
fuel Int
count Graph
g Int
source Int
target = do
    Graph -> Int -> Either MathError ()
checkVertex Graph
g Int
source
    Graph -> Int -> Either MathError ()
checkVertex Graph
g Int
target
    Graph -> Either MathError ()
nonnegative Graph
g
    if Int
fuel Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        then MathError -> Either MathError [(Double, [Int], [Int])]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative path count/budget")
        else Int
-> [(Double, [Int], [Int])]
-> [(Double, [Int], [Int])]
-> Either MathError [(Double, [Int], [Int])]
forall {t}.
(Ord t, Num t) =>
t
-> [(Double, [Int], [Int])]
-> [(Double, [Int], [Int])]
-> Either MathError [(Double, [Int], [Int])]
go Int
fuel [(Double
0, [Int
source], [])] []
  where
    go :: t
-> [(Double, [Int], [Int])]
-> [(Double, [Int], [Int])]
-> Either MathError [(Double, [Int], [Int])]
go t
_ [(Double, [Int], [Int])]
_ [(Double, [Int], [Int])]
out | [(Double, [Int], [Int])] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Double, [Int], [Int])]
out Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
count = [(Double, [Int], [Int])]
-> Either MathError [(Double, [Int], [Int])]
forall a b. b -> Either a b
Right [(Double, [Int], [Int])]
out
    go t
_ [] [(Double, [Int], [Int])]
out = [(Double, [Int], [Int])]
-> Either MathError [(Double, [Int], [Int])]
forall a b. b -> Either a b
Right [(Double, [Int], [Int])]
out
    go t
b [(Double, [Int], [Int])]
_ [(Double, [Int], [Int])]
_ | t
b t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 = MathError -> Either MathError [(Double, [Int], [Int])]
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"K shortest simple paths" Int
fuel)
    go t
b [(Double, [Int], [Int])]
queue [(Double, [Int], [Int])]
out = case ((Double, [Int], [Int]) -> (Double, [Int], [Int]))
-> [(Double, [Int], [Int])] -> [(Double, [Int], [Int])]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Double, [Int], [Int]) -> (Double, [Int], [Int])
forall a. a -> a
id [(Double, [Int], [Int])]
queue of
        [] -> [(Double, [Int], [Int])]
-> Either MathError [(Double, [Int], [Int])]
forall a b. b -> Either a b
Right [(Double, [Int], [Int])]
out
        ((Double
cost, [Int]
path, [Int]
ids) : [(Double, [Int], [Int])]
rest) ->
            let v :: Int
v = [Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
path
             in if Bool -> Bool
not (Double -> Bool
finite Double
cost)
                    then MathError -> Either MathError [(Double, [Int], [Int])]
forall a b. a -> Either a b
Left MathError
NonFiniteResult
                    else
                        if Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
target
                            then t
-> [(Double, [Int], [Int])]
-> [(Double, [Int], [Int])]
-> Either MathError [(Double, [Int], [Int])]
go (t
b t -> t -> t
forall a. Num a => a -> a -> a
- t
1) [(Double, [Int], [Int])]
rest ([(Double, [Int], [Int])]
out [(Double, [Int], [Int])]
-> [(Double, [Int], [Int])] -> [(Double, [Int], [Int])]
forall a. [a] -> [a] -> [a]
++ [(Double
cost, [Int]
path, [Int]
ids)])
                            else do
                                added <-
                                    ((Int, Double, Int) -> Either MathError (Double, [Int], [Int]))
-> [(Int, Double, Int)]
-> Either MathError [(Double, [Int], [Int])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
                                        (\(Int
u, Double
w, Int
i) -> do candidate <- Double -> Double -> Either MathError Double
checkedSum Double
cost Double
w; pure (candidate, path ++ [u], ids ++ [i]))
                                        [(Int
u, Double
w, Int
i) | (Int
u, Double
w, Int
i) <- Graph -> Int -> [(Int, Double, Int)]
neighbors Graph
g Int
v, Int
u Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
path]
                                go (b - 1) (rest ++ added) out

{- | Given directed acyclic graph then source, return maximum path-weight distance to each
vertex, Nothing for unreachable vertices. Invalid source, cycles/undirected input or
non-finite arithmetic fails.
-}
longestDAGPaths :: Graph -> Int -> Either MathError [Maybe Double]
longestDAGPaths :: Graph -> Int -> Either MathError [Maybe Double]
longestDAGPaths Graph
g Int
source = do
    Graph -> Int -> Either MathError ()
checkVertex Graph
g Int
source
    order <- Graph -> Either MathError [Int]
topologicalSort Graph
g
    result <-
        foldM
            ( \Distances
m Int
v -> case Int -> Distances -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
v Distances
m of
                Maybe Double
Nothing -> Distances -> Either MathError Distances
forall a b. b -> Either a b
Right Distances
m
                Just Double
d ->
                    (Distances -> (Int, Double, Int) -> Either MathError Distances)
-> Distances -> [(Int, Double, Int)] -> Either MathError Distances
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
                        (\Distances
z (Int
u, Double
w, Int
_) -> do candidate <- Double -> Double -> Either MathError Double
checkedSum Double
d Double
w; pure (M.insertWith max u candidate z))
                        Distances
m
                        (Graph -> Int -> [(Int, Double, Int)]
neighbors Graph
g Int
v)
            )
            (M.singleton source 0)
            order
    finish g result