-- SPDX-License-Identifier: MPL-2.0

{- | Exact integral network flow and assignment. Capacities and costs use
'Integer'; residual arcs preserve parallel and antiparallel edge identity.
Each max-flow algorithm has a caller-supplied augmentation/push/relabel budget.
-}
module HMath.Graph.Flow (
    Network,
    NetworkEdge (..),
    network,
    networkEdges,
    FlowResult (..),
    fordFulkerson,
    edmondsKarp,
    dinic,
    pushRelabel,
    minimumCut,
    minimumCostMaximumFlow,
    hopcroftKarp,
    hungarian,
) where

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

{- | Directed integral-capacity/unit-cost edge with stable identity. Capacities are
nonnegative; costs may be negative.
-}
data NetworkEdge = NetworkEdge
    { NetworkEdge -> Int
networkEdgeId :: Int
    -- ^ Stable original edge identity in network input order.
    , NetworkEdge -> Int
networkFrom :: Int
    -- ^ Source vertex of this directed network edge.
    , NetworkEdge -> Int
networkTo :: Int
    -- ^ Destination vertex of this directed network edge.
    , NetworkEdge -> Integer
networkCapacity :: Integer
    -- ^ Exact nonnegative Integer capacity of this edge.
    , NetworkEdge -> Integer
networkCost :: Integer
    -- ^ Exact Integer cost per unit flow; may be negative.
    }
    deriving (NetworkEdge -> NetworkEdge -> Bool
(NetworkEdge -> NetworkEdge -> Bool)
-> (NetworkEdge -> NetworkEdge -> Bool) -> Eq NetworkEdge
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NetworkEdge -> NetworkEdge -> Bool
== :: NetworkEdge -> NetworkEdge -> Bool
$c/= :: NetworkEdge -> NetworkEdge -> Bool
/= :: NetworkEdge -> NetworkEdge -> Bool
Eq, Int -> NetworkEdge -> ShowS
[NetworkEdge] -> ShowS
NetworkEdge -> String
(Int -> NetworkEdge -> ShowS)
-> (NetworkEdge -> String)
-> ([NetworkEdge] -> ShowS)
-> Show NetworkEdge
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NetworkEdge -> ShowS
showsPrec :: Int -> NetworkEdge -> ShowS
$cshow :: NetworkEdge -> String
show :: NetworkEdge -> String
$cshowList :: [NetworkEdge] -> ShowS
showList :: [NetworkEdge] -> ShowS
Show)

{- | Opaque exact network-flow input with stable parallel/antiparallel edges, constructed
with network.
-}
data Network = Network Int [NetworkEdge] deriving (Network -> Network -> Bool
(Network -> Network -> Bool)
-> (Network -> Network -> Bool) -> Eq Network
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Network -> Network -> Bool
== :: Network -> Network -> Bool
$c/= :: Network -> Network -> Bool
/= :: Network -> Network -> Bool
Eq, Int -> Network -> ShowS
[Network] -> ShowS
Network -> String
(Int -> Network -> ShowS)
-> (Network -> String) -> ([Network] -> ShowS) -> Show Network
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Network -> ShowS
showsPrec :: Int -> Network -> ShowS
$cshow :: Network -> String
show :: Network -> String
$cshowList :: [Network] -> ShowS
showList :: [Network] -> ShowS
Show)

{- | Given vertex count then (from,to,capacity,cost) edges, construct an exact integral
network with stable edge IDs. Require \<=512 vertices, \<=65536 edges, valid endpoints
and nonnegative capacities. Costs may be negative.
-}
network :: Int -> [(Int, Int, Integer, Integer)] -> Either MathError Network
network :: Int -> [(Int, Int, Integer, Integer)] -> Either MathError Network
network Int
n [(Int, Int, Integer, Integer)]
es
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
512
        Bool -> Bool -> Bool
|| [(Int, Int, Integer, Integer)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Int, Integer, Integer)]
es Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
65536
        Bool -> Bool -> Bool
|| ((Int, Int, Integer, Integer) -> Bool)
-> [(Int, Int, Integer, Integer)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Int
u, Int
v, Integer
c, Integer
_) -> Int
u Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
u Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n Bool -> Bool -> Bool
|| Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n Bool -> Bool -> Bool
|| Integer
c Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0) [(Int, Int, Integer, Integer)]
es =
        MathError -> Either MathError Network
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid network size/endpoints/capacity")
    | Bool
otherwise = Network -> Either MathError Network
forall a b. b -> Either a b
Right (Int -> [NetworkEdge] -> Network
Network Int
n [Int -> Int -> Int -> Integer -> Integer -> NetworkEdge
NetworkEdge Int
i Int
u Int
v Integer
c Integer
w | (Int
i, (Int
u, Int
v, Integer
c, Integer
w)) <- [Int]
-> [(Int, Int, Integer, Integer)]
-> [(Int, (Int, Int, Integer, Integer))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [(Int, Int, Integer, Integer)]
es])

-- | Return original network edges in input/ID order, retaining capacities and costs.
networkEdges :: Network -> [NetworkEdge]
networkEdges :: Network -> [NetworkEdge]
networkEdges (Network Int
_ [NetworkEdge]
es) = [NetworkEdge]
es

{- | Exact net flow,total edge cost,and original-edge-order flows. A maximum-flow result
alone does not certify minimum cost.
-}
data FlowResult = FlowResult
    { FlowResult -> Integer
flowValue :: Integer
    -- ^ Exact net flow leaving the chosen source.
    , FlowResult -> [Integer]
edgeFlows :: [Integer]
    -- ^ Exact edge flows in original network edge order, preserving parallel and antiparallel identity.
    , FlowResult -> Integer
flowCost :: Integer
    -- ^ Sum of edge flow times edge unit cost; maximum-flow methods do not necessarily minimize this cost.
    }
    deriving (FlowResult -> FlowResult -> Bool
(FlowResult -> FlowResult -> Bool)
-> (FlowResult -> FlowResult -> Bool) -> Eq FlowResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: FlowResult -> FlowResult -> Bool
== :: FlowResult -> FlowResult -> Bool
$c/= :: FlowResult -> FlowResult -> Bool
/= :: FlowResult -> FlowResult -> Bool
Eq, Int -> FlowResult -> ShowS
[FlowResult] -> ShowS
FlowResult -> String
(Int -> FlowResult -> ShowS)
-> (FlowResult -> String)
-> ([FlowResult] -> ShowS)
-> Show FlowResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> FlowResult -> ShowS
showsPrec :: Int -> FlowResult -> ShowS
$cshow :: FlowResult -> String
show :: FlowResult -> String
$cshowList :: [FlowResult] -> ShowS
showList :: [FlowResult] -> ShowS
Show)

type Residual = M.Map Int Integer
arcData :: Network -> Int -> (Int, Int, Integer)
arcData :: Network -> Int -> (Int, Int, Integer)
arcData (Network Int
_ [NetworkEdge]
es) Int
i =
    let e :: NetworkEdge
e = [NetworkEdge]
es [NetworkEdge] -> Int -> NetworkEdge
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
     in if Int -> Bool
forall a. Integral a => a -> Bool
even Int
i
            then (NetworkEdge -> Int
networkFrom NetworkEdge
e, NetworkEdge -> Int
networkTo NetworkEdge
e, NetworkEdge -> Integer
networkCost NetworkEdge
e)
            else (NetworkEdge -> Int
networkTo NetworkEdge
e, NetworkEdge -> Int
networkFrom NetworkEdge
e, -NetworkEdge -> Integer
networkCost NetworkEdge
e)
outgoing :: Network -> Residual -> Int -> [Int]
outgoing :: Network -> Residual -> Int -> [Int]
outgoing net :: Network
net@(Network Int
_ [NetworkEdge]
es) Residual
residual Int
u =
    [ Int
i
    | Int
i <- [Int
0 .. Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* [NetworkEdge] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [NetworkEdge]
es Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    , let (Int
a, Int
b, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i
    , Int
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
u
    , Int
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
b
    , Integer -> Int -> Residual -> Integer
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Integer
0 Int
i Residual
residual Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0
    ]
initial :: Network -> Residual
initial :: Network -> Residual
initial (Network Int
_ [NetworkEdge]
es) =
    [(Int, Integer)] -> Residual
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
        ([[(Int, Integer)]] -> [(Int, Integer)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* NetworkEdge -> Int
networkEdgeId NetworkEdge
e, NetworkEdge -> Integer
networkCapacity NetworkEdge
e), (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* NetworkEdge -> Int
networkEdgeId NetworkEdge
e Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Integer
0)] | NetworkEdge
e <- [NetworkEdge]
es])
augment :: Integer -> [Int] -> Residual -> Residual
augment :: Integer -> [Int] -> Residual -> Residual
augment Integer
amount [Int]
path Residual
r = (Residual -> Int -> Residual) -> Residual -> [Int] -> Residual
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\Residual
m Int
i -> (Integer -> Integer) -> Int -> Residual -> Residual
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
subtract Integer
amount) Int
i ((Integer -> Integer) -> Int -> Residual -> Residual
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
amount) (Int -> Int
forall {a}. Integral a => a -> a
xorOne Int
i) Residual
m)) Residual
r [Int]
path
  where
    xorOne :: a -> a
xorOne a
i = if a -> Bool
forall a. Integral a => a -> Bool
even a
i then a
i a -> a -> a
forall a. Num a => a -> a -> a
+ a
1 else a
i a -> a -> a
forall a. Num a => a -> a -> a
- a
1
result :: Network -> Int -> Residual -> FlowResult
result :: Network -> Int -> Residual -> FlowResult
result (Network Int
_ [NetworkEdge]
es) Int
source Residual
r = Integer -> [Integer] -> Integer -> FlowResult
FlowResult Integer
value [Integer]
flows Integer
cost
  where
    flows :: [Integer]
flows = [NetworkEdge -> Integer
networkCapacity NetworkEdge
e Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer -> Int -> Residual -> Integer
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Integer
0 (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* NetworkEdge -> Int
networkEdgeId NetworkEdge
e) Residual
r | NetworkEdge
e <- [NetworkEdge]
es]
    value :: Integer
value =
        [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
f | (NetworkEdge
e, Integer
f) <- [NetworkEdge] -> [Integer] -> [(NetworkEdge, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip [NetworkEdge]
es [Integer]
flows, NetworkEdge -> Int
networkFrom NetworkEdge
e Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
source]
            Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
f | (NetworkEdge
e, Integer
f) <- [NetworkEdge] -> [Integer] -> [(NetworkEdge, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip [NetworkEdge]
es [Integer]
flows, NetworkEdge -> Int
networkTo NetworkEdge
e Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
source]
    cost :: Integer
cost = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
f Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* NetworkEdge -> Integer
networkCost NetworkEdge
e | (NetworkEdge
e, Integer
f) <- [NetworkEdge] -> [Integer] -> [(NetworkEdge, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip [NetworkEdge]
es [Integer]
flows]
valid :: Int -> Network -> Int -> Int -> Either MathError ()
valid :: Int -> Network -> Int -> Int -> Either MathError ()
valid Int
budget (Network Int
n [NetworkEdge]
_) Int
s Int
t
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n Bool -> Bool -> Bool
|| Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n Bool -> Bool -> Bool
|| Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
t =
        MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid flow budget/source/sink")
    | Bool
otherwise = () -> Either MathError ()
forall a b. b -> Either a b
Right ()

findPath :: Bool -> Network -> Residual -> Int -> Int -> Maybe [Int]
findPath :: Bool -> Network -> Residual -> Int -> Int -> Maybe [Int]
findPath Bool
breadth Network
net Residual
r Int
s Int
t = Set Int -> [(Int, [Int])] -> Maybe [Int]
go Set Int
forall a. Set a
S.empty [(Int
s, [])]
  where
    go :: Set Int -> [(Int, [Int])] -> Maybe [Int]
go Set Int
_ [] = Maybe [Int]
forall a. Maybe a
Nothing
    go Set Int
seen ((Int
v, [Int]
path) : [(Int, [Int])]
rest)
        | Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
t = [Int] -> Maybe [Int]
forall a. a -> Maybe a
Just [Int]
path
        | Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Int
v Set Int
seen = Set Int -> [(Int, [Int])] -> Maybe [Int]
go Set Int
seen [(Int, [Int])]
rest
        | Bool
otherwise =
            let next :: [(Int, [Int])]
next = [(Int
b, [Int]
path [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
i]) | Int
i <- Network -> Residual -> Int -> [Int]
outgoing Network
net Residual
r Int
v, let (Int
_, Int
b, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i]
             in Set Int -> [(Int, [Int])] -> Maybe [Int]
go (Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
S.insert Int
v Set Int
seen) (if Bool
breadth then [(Int, [Int])]
rest [(Int, [Int])] -> [(Int, [Int])] -> [(Int, [Int])]
forall a. [a] -> [a] -> [a]
++ [(Int, [Int])]
next else [(Int, [Int])]
next [(Int, [Int])] -> [(Int, [Int])] -> [(Int, [Int])]
forall a. [a] -> [a] -> [a]
++ [(Int, [Int])]
rest)
augmenting :: Bool -> Int -> Network -> Int -> Int -> Either MathError FlowResult
augmenting :: Bool -> Int -> Network -> Int -> Int -> Either MathError FlowResult
augmenting Bool
breadth Int
b Network
net Int
s Int
t = do Int -> Network -> Int -> Int -> Either MathError ()
valid Int
b Network
net Int
s Int
t; Int -> Residual -> Either MathError FlowResult
forall {t}.
(Ord t, Num t) =>
t -> Residual -> Either MathError FlowResult
go Int
b (Network -> Residual
initial Network
net)
  where
    go :: t -> Residual -> Either MathError FlowResult
go t
fuel Residual
r = case Bool -> Network -> Residual -> Int -> Int -> Maybe [Int]
findPath Bool
breadth Network
net Residual
r Int
s Int
t of
        Maybe [Int]
Nothing -> FlowResult -> Either MathError FlowResult
forall a b. b -> Either a b
Right (Network -> Int -> Residual -> FlowResult
result Network
net Int
s Residual
r)
        Just [Int]
path
            | t
fuel t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 -> MathError -> Either MathError FlowResult
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"augmenting maximum flow" Int
b)
            | Bool
otherwise -> t -> Residual -> Either MathError FlowResult
go (t
fuel t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Integer -> [Int] -> Residual -> Residual
augment ([Integer] -> Integer
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Residual
r Residual -> Int -> Integer
forall k a. Ord k => Map k a -> k -> a
M.! Int
i | Int
i <- [Int]
path]) [Int]
path Residual
r)
fordFulkerson, edmondsKarp :: Int -> Network -> Int -> Int -> Either MathError FlowResult

{- | Given augmentation budget,network,source,sink, compute exact maximum flow using
depth-first augmenting paths. Invalid/non-distinct endpoints or negative budget fails;
exhausted budget returns NonConvergence. Result retains per-edge flows and total cost.
-}
fordFulkerson :: Int -> Network -> Int -> Int -> Either MathError FlowResult
fordFulkerson = Bool -> Int -> Network -> Int -> Int -> Either MathError FlowResult
augmenting Bool
False

{- | Given augmentation budget,network,source,sink, compute exact maximum flow using
breadth-first shortest augmenting paths. Invalid input or exhausted budget fails;
parallel/antiparallel edge flows retain identity.
-}
edmondsKarp :: Int -> Network -> Int -> Int -> Either MathError FlowResult
edmondsKarp = Bool -> Int -> Network -> Int -> Int -> Either MathError FlowResult
augmenting Bool
True

{- | Given work budget,network,source,sink, compute maximum flow with level graphs and
blocking-flow augmentation. Invalid source/sink/budget or exhausted work returns Left
MathError; result includes flows in original edge order.
-}
dinic :: Int -> Network -> Int -> Int -> Either MathError FlowResult
dinic :: Int -> Network -> Int -> Int -> Either MathError FlowResult
dinic Int
b Network
net Int
s Int
t = do Int -> Network -> Int -> Int -> Either MathError ()
valid Int
b Network
net Int
s Int
t; Int -> Residual -> Either MathError FlowResult
forall {t}.
(Ord t, Num t) =>
t -> Residual -> Either MathError FlowResult
phase Int
b (Network -> Residual
initial Network
net)
  where
    levels :: Residual -> Map Int Int
levels Residual
r = Map Int Int -> [Int] -> Map Int Int
forall {a}. Num a => Map Int a -> [Int] -> Map Int a
go (Int -> Int -> Map Int Int
forall k a. k -> a -> Map k a
M.singleton Int
s (Int
0 :: Int)) [Int
s]
      where
        go :: Map Int a -> [Int] -> Map Int a
go Map Int a
dist [] = Map Int a
dist
        go Map Int a
dist (Int
v : [Int]
vs) =
            let fresh :: [Int]
fresh = [Int
u | Int
i <- Network -> Residual -> Int -> [Int]
outgoing Network
net Residual
r Int
v, let (Int
_, Int
u, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i, Int -> Map Int a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.notMember Int
u Map Int a
dist]
                next :: Map Int a
next = (Map Int a -> Int -> Map Int a) -> Map Int a -> [Int] -> Map Int a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\Map Int a
m Int
u -> Int -> a -> Map Int a -> Map Int a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
u (Map Int a
dist Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
v a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) Map Int a
m) Map Int a
dist [Int]
fresh
             in Map Int a -> [Int] -> Map Int a
go Map Int a
next ([Int]
vs [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Set Int -> [Int]
forall a. Set a -> [a]
S.toList ([Int] -> Set Int
forall a. Ord a => [a] -> Set a
S.fromList [Int]
fresh))
    layered :: Residual -> Map Int a -> Int -> Maybe [Int]
layered Residual
r Map Int a
dist Int
v
        | Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
t = [Int] -> Maybe [Int]
forall a. a -> Maybe a
Just []
        | Bool
otherwise =
            [Maybe [Int]] -> Maybe [Int]
forall {a}. [Maybe a] -> Maybe a
first
                [ (Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:) ([Int] -> [Int]) -> Maybe [Int] -> Maybe [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Residual -> Map Int a -> Int -> Maybe [Int]
layered Residual
r Map Int a
dist Int
u
                | Int
i <- Network -> Residual -> Int -> [Int]
outgoing Network
net Residual
r Int
v
                , let (Int
_, Int
u, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i
                , Int -> Map Int a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
u Map Int a
dist Maybe a -> Maybe a -> Bool
forall a. Eq a => a -> a -> Bool
== ((a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (a -> a) -> Maybe a -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Map Int a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
v Map Int a
dist)
                ]
    first :: [Maybe a] -> Maybe a
first [] = Maybe a
forall a. Maybe a
Nothing; first (Just a
x : [Maybe a]
_) = a -> Maybe a
forall a. a -> Maybe a
Just a
x; first (Maybe a
Nothing : [Maybe a]
xs) = [Maybe a] -> Maybe a
first [Maybe a]
xs
    phase :: a -> Residual -> Either MathError FlowResult
phase a
fuel Residual
r = let dist :: Map Int Int
dist = Residual -> Map Int Int
levels Residual
r in if Int -> Map Int Int -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.notMember Int
t Map Int Int
dist then FlowResult -> Either MathError FlowResult
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Network -> Int -> Residual -> FlowResult
result Network
net Int
s Residual
r) else a -> Residual -> Map Int Int -> Either MathError FlowResult
block a
fuel Residual
r Map Int Int
dist
    block :: a -> Residual -> Map Int Int -> Either MathError FlowResult
block a
fuel Residual
r Map Int Int
dist = case Residual -> Map Int Int -> Int -> Maybe [Int]
forall {a}.
(Eq a, Num a) =>
Residual -> Map Int a -> Int -> Maybe [Int]
layered Residual
r Map Int Int
dist Int
s of
        Maybe [Int]
Nothing -> a -> Residual -> Either MathError FlowResult
phase a
fuel Residual
r
        Just [Int]
path
            | a
fuel a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 -> MathError -> Either MathError FlowResult
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Dinic maximum flow" Int
b)
            | Bool
otherwise -> a -> Residual -> Map Int Int -> Either MathError FlowResult
block (a
fuel a -> a -> a
forall a. Num a => a -> a -> a
- a
1) (Integer -> [Int] -> Residual -> Residual
augment ([Integer] -> Integer
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Residual
r Residual -> Int -> Integer
forall k a. Ord k => Map k a -> k -> a
M.! Int
i | Int
i <- [Int]
path]) [Int]
path Residual
r) Map Int Int
dist

{- | Given push/relabel work budget,network,source,sink, compute exact maximum flow by
preflow and relabel operations. Reject invalid endpoints/budget; work exhaustion
returns NonConvergence.
-}
pushRelabel :: Int -> Network -> Int -> Int -> Either MathError FlowResult
pushRelabel :: Int -> Network -> Int -> Int -> Either MathError FlowResult
pushRelabel Int
b net :: Network
net@(Network Int
n [NetworkEdge]
_) Int
s Int
t = do
    Int -> Network -> Int -> Int -> Either MathError ()
valid Int
b Network
net Int
s Int
t
    let r0 :: Residual
r0 = Network -> Residual
initial Network
net
        pre :: [Int]
pre = Network -> Residual -> Int -> [Int]
outgoing Network
net Residual
r0 Int
s
        (Residual
r, Residual
excess) =
            ((Residual, Residual) -> Int -> (Residual, Residual))
-> (Residual, Residual) -> [Int] -> (Residual, Residual)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl'
                ( \(Residual
res, Residual
ex) Int
i ->
                    let amount :: Integer
amount = Residual
res Residual -> Int -> Integer
forall k a. Ord k => Map k a -> k -> a
M.! Int
i; (Int
_, Int
u, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i
                     in (Integer -> [Int] -> Residual -> Residual
augment Integer
amount [Int
i] Residual
res, (Integer -> Integer -> Integer)
-> Int -> Integer -> Residual -> Residual
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) Int
u Integer
amount ((Integer -> Integer -> Integer)
-> Int -> Integer -> Residual -> Residual
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) Int
s (-Integer
amount) Residual
ex))
                )
                (Residual
r0, Residual
forall k a. Map k a
M.empty)
                [Int]
pre
    Int
-> Residual
-> Residual
-> Map Int Int
-> Either MathError FlowResult
forall {t} {a}.
(Num t, Num a, Ord t, Ord a) =>
t
-> Residual -> Residual -> Map Int a -> Either MathError FlowResult
go Int
b Residual
r Residual
excess (Int -> Int -> Map Int Int -> Map Int Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
s Int
n ([(Int, Int)] -> Map Int Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
v, Int
0) | Int
v <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]))
  where
    go :: t
-> Residual -> Residual -> Map Int a -> Either MathError FlowResult
go t
fuel Residual
r Residual
excess Map Int a
heights = case [Int
v | Int
v <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
s, Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
t, Integer -> Int -> Residual -> Integer
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Integer
0 Int
v Residual
excess Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0] of
        [] -> FlowResult -> Either MathError FlowResult
forall a b. b -> Either a b
Right (Network -> Int -> Residual -> FlowResult
result Network
net Int
s Residual
r)
        [Int]
_ | t
fuel t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 -> MathError -> Either MathError FlowResult
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"push-relabel maximum flow" Int
b)
        (Int
v : [Int]
_) -> case [Int
i | Int
i <- Network -> Residual -> Int -> [Int]
outgoing Network
net Residual
r Int
v, let (Int
_, Int
u, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i, Map Int a
heights Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Map Int a
heights Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
u a -> a -> a
forall a. Num a => a -> a -> a
+ a
1] of
            (Int
i : [Int]
_) ->
                let (Int
_, Int
u, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i; amount :: Integer
amount = Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min (Residual
excess Residual -> Int -> Integer
forall k a. Ord k => Map k a -> k -> a
M.! Int
v) (Residual
r Residual -> Int -> Integer
forall k a. Ord k => Map k a -> k -> a
M.! Int
i)
                 in t
-> Residual -> Residual -> Map Int a -> Either MathError FlowResult
go
                        (t
fuel t -> t -> t
forall a. Num a => a -> a -> a
- t
1)
                        (Integer -> [Int] -> Residual -> Residual
augment Integer
amount [Int
i] Residual
r)
                        ((Integer -> Integer -> Integer)
-> Int -> Integer -> Residual -> Residual
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) Int
u Integer
amount ((Integer -> Integer) -> Int -> Residual -> Residual
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
subtract Integer
amount) Int
v Residual
excess))
                        Map Int a
heights
            [] -> case [Map Int a
heights Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
u | Int
i <- Network -> Residual -> Int -> [Int]
outgoing Network
net Residual
r Int
v, let (Int
_, Int
u, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i] of
                [] -> MathError -> Either MathError FlowResult
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"invalid residual preflow")
                [a]
hs -> t
-> Residual -> Residual -> Map Int a -> Either MathError FlowResult
go (t
fuel t -> t -> t
forall a. Num a => a -> a -> a
- t
1) Residual
r Residual
excess (Int -> a -> Map Int a -> Map Int a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
v ([a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [a]
hs a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) Map Int a
heights)

-- | Source-side vertices and cut capacity, using Edmonds--Karp residual reach.
minimumCut :: Int -> Network -> Int -> Int -> Either MathError ([Int], Integer)
minimumCut :: Int -> Network -> Int -> Int -> Either MathError ([Int], Integer)
minimumCut Int
b net :: Network
net@(Network Int
_ [NetworkEdge]
es) Int
s Int
t = do
    answer <- Int -> Network -> Int -> Int -> Either MathError FlowResult
edmondsKarp Int
b Network
net Int
s Int
t
    let r = Network -> [Integer] -> Residual
fromFlows Network
net (FlowResult -> [Integer]
edgeFlows FlowResult
answer)
        reach Set Int
seen [] = Set Int
seen
        reach Set Int
seen (Int
v : [Int]
vs)
            | Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Int
v Set Int
seen = Set Int -> [Int] -> Set Int
reach Set Int
seen [Int]
vs
            | Bool
otherwise =
                Set Int -> [Int] -> Set Int
reach (Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
S.insert Int
v Set Int
seen) ([Int
u | Int
i <- Network -> Residual -> Int -> [Int]
outgoing Network
net Residual
r Int
v, let { (Int
_, Int
u, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i }] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
vs)
        side = Set Int -> [Int] -> Set Int
reach Set Int
forall a. Set a
S.empty [Int
s]
    pure
        ( S.toAscList side
        , sum [networkCapacity e | e <- es, S.member (networkFrom e) side, S.notMember (networkTo e) side]
        )

fromFlows :: Network -> [Integer] -> Residual
fromFlows :: Network -> [Integer] -> Residual
fromFlows (Network Int
_ [NetworkEdge]
es) [Integer]
fs =
    [(Int, Integer)] -> Residual
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
        ( [[(Int, Integer)]] -> [(Int, Integer)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
            [[(Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* NetworkEdge -> Int
networkEdgeId NetworkEdge
e, NetworkEdge -> Integer
networkCapacity NetworkEdge
e Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
f), (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* NetworkEdge -> Int
networkEdgeId NetworkEdge
e Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Integer
f)] | (NetworkEdge
e, Integer
f) <- [NetworkEdge] -> [Integer] -> [(NetworkEdge, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip [NetworkEdge]
es [Integer]
fs]
        )

{- | Maximum flow followed by exact negative-residual-cycle cancellation. This
supports negative costs (including disconnected cycles) and finds the minimum
cost among all maximum flows. The same budget bounds each phase separately.
-}
minimumCostMaximumFlow :: Int -> Network -> Int -> Int -> Either MathError FlowResult
minimumCostMaximumFlow :: Int -> Network -> Int -> Int -> Either MathError FlowResult
minimumCostMaximumFlow Int
b net :: Network
net@(Network Int
n [NetworkEdge]
es) Int
s Int
t = do
    start <- Int -> Network -> Int -> Int -> Either MathError FlowResult
dinic Int
b Network
net Int
s Int
t
    cancel b (fromFlows net (edgeFlows start))
  where
    cancel :: t -> Residual -> Either MathError FlowResult
cancel t
fuel Residual
r = case Residual -> Maybe [Int]
forall {a}. (Ord a, Num a) => Map Int a -> Maybe [Int]
negativeCycle Residual
r of
        Maybe [Int]
Nothing -> FlowResult -> Either MathError FlowResult
forall a b. b -> Either a b
Right (Network -> Int -> Residual -> FlowResult
result Network
net Int
s Residual
r)
        Just [Int]
path
            | t
fuel t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 -> MathError -> Either MathError FlowResult
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"minimum-cost cycle cancellation" Int
b)
            | Bool
otherwise -> t -> Residual -> Either MathError FlowResult
cancel (t
fuel t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Integer -> [Int] -> Residual -> Residual
augment ([Integer] -> Integer
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Residual
r Residual -> Int -> Integer
forall k a. Ord k => Map k a -> k -> a
M.! Int
i | Int
i <- [Int]
path]) [Int]
path Residual
r)
    negativeCycle :: Map Int a -> Maybe [Int]
negativeCycle Map Int a
r =
        let active :: [Int]
active = [Int
i | Int
i <- [Int
0 .. Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* [NetworkEdge] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [NetworkEdge]
es Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Map Int a
r Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0]
            iteration :: (Residual, Map Int Int, c) -> (Residual, Map Int Int, Maybe Int)
iteration (Residual
dist, Map Int Int
parent, c
_) =
                ((Residual, Map Int Int, Maybe Int)
 -> Int -> (Residual, Map Int Int, Maybe Int))
-> (Residual, Map Int Int, Maybe Int)
-> [Int]
-> (Residual, Map Int Int, Maybe Int)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl'
                    ( \(Residual
d, Map Int Int
p, Maybe Int
lastChanged) Int
i ->
                        let (Int
u, Int
v, Integer
c) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i; candidate :: Integer
candidate = Residual
d Residual -> Int -> Integer
forall k a. Ord k => Map k a -> k -> a
M.! Int
u Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
c
                         in if Integer
candidate Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Residual
d Residual -> Int -> Integer
forall k a. Ord k => Map k a -> k -> a
M.! Int
v
                                then (Int -> Integer -> Residual -> Residual
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
v Integer
candidate Residual
d, Int -> Int -> Map Int Int -> Map Int Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
v Int
i Map Int Int
p, Int -> Maybe Int
forall a. a -> Maybe a
Just Int
v)
                                else (Residual
d, Map Int Int
p, Maybe Int
lastChanged)
                    )
                    (Residual
dist, Map Int Int
parent, Maybe Int
forall a. Maybe a
Nothing)
                    [Int]
active
            (Residual
_, Map Int Int
parents, Maybe Int
lastVertex) = ((Residual, Map Int Int, Maybe Int)
 -> (Residual, Map Int Int, Maybe Int))
-> (Residual, Map Int Int, Maybe Int)
-> [(Residual, Map Int Int, Maybe Int)]
forall a. (a -> a) -> a -> [a]
iterate (Residual, Map Int Int, Maybe Int)
-> (Residual, Map Int Int, Maybe Int)
forall {c}.
(Residual, Map Int Int, c) -> (Residual, Map Int Int, Maybe Int)
iteration ([(Int, Integer)] -> Residual
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
v, Integer
0 :: Integer) | Int
v <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]], Map Int Int
forall k a. Map k a
M.empty, Maybe Int
forall a. Maybe a
Nothing) [(Residual, Map Int Int, Maybe Int)]
-> Int -> (Residual, Map Int Int, Maybe Int)
forall a. HasCallStack => [a] -> Int -> a
!! Int
n
            previous :: Int -> Int
previous Int
v = let i :: Int
i = Map Int Int
parents Map Int Int -> Int -> Int
forall k a. Ord k => Map k a -> k -> a
M.! Int
v; (Int
u, Int
_, Integer
_) = Network -> Int -> (Int, Int, Integer)
arcData Network
net Int
i in Int
u
            build :: Int -> Int -> [Int] -> [Int]
build Int
start Int
v [Int]
acc = let i :: Int
i = Map Int Int
parents Map Int Int -> Int -> Int
forall k a. Ord k => Map k a -> k -> a
M.! Int
v; u :: Int
u = Int -> Int
previous Int
v in if Int
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
start then Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
acc else Int -> Int -> [Int] -> [Int]
build Int
start Int
u (Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
acc)
         in case Maybe Int
lastVertex of
                Maybe Int
Nothing -> Maybe [Int]
forall a. Maybe a
Nothing
                Just Int
v -> let start :: Int
start = (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate Int -> Int
previous Int
v [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
n in [Int] -> Maybe [Int]
forall a. a -> Maybe a
Just (Int -> Int -> [Int] -> [Int]
build Int
start Int
start [])

{- | Hopcroft--Karp layered augmenting paths. Returns @(left,right)@ matching;
parallel edges are coalesced and weights are ignored.
-}
hopcroftKarp :: BipartiteGraph -> [(Int, Int)]
hopcroftKarp :: BipartiteGraph -> [(Int, Int)]
hopcroftKarp BipartiteGraph
bip = [(Int, Int)] -> [(Int, Int)]
forall a. Ord a => [a] -> [a]
sort (Map Int Int -> [(Int, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Int Int -> Map Int Int -> Map Int Int
go Map Int Int
forall k a. Map k a
M.empty Map Int Int
forall k a. Map k a
M.empty))
  where
    (Int
left, Int
_) = BipartiteGraph -> (Int, Int)
bipartiteParts BipartiteGraph
bip
    g :: Graph
g = BipartiteGraph -> Graph
bipartiteBase BipartiteGraph
bip
    adj :: Int -> [Int]
adj Int
u = Set Int -> [Int]
forall a. Set a -> [a]
S.toAscList ([Int] -> Set Int
forall a. Ord a => [a] -> Set a
S.fromList [Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
left | (Int
v, Double
_, Int
_) <- Graph -> Int -> [(Int, Double, Int)]
neighbors Graph
g Int
u])
    bfsLayer :: Map Int a -> Map Int Int -> (Map Int Int, Maybe Int)
bfsLayer Map Int a
matching Map Int Int
reverseMatching =
        Map Int Int -> [Int] -> Maybe Int -> (Map Int Int, Maybe Int)
forall {a}.
(Num a, Ord a) =>
Map Int a -> [Int] -> Maybe a -> (Map Int a, Maybe a)
walk
            ([(Int, Int)] -> Map Int Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
u, Int
0 :: Int) | Int
u <- [Int
0 .. Int
left Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int -> Map Int a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.notMember Int
u Map Int a
matching])
            [Int
u | Int
u <- [Int
0 .. Int
left Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int -> Map Int a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.notMember Int
u Map Int a
matching]
            Maybe Int
forall a. Maybe a
Nothing
      where
        walk :: Map Int a -> [Int] -> Maybe a -> (Map Int a, Maybe a)
walk Map Int a
dist [] Maybe a
shortest = (Map Int a
dist, Maybe a
shortest)
        walk Map Int a
dist (Int
u : [Int]
us) Maybe a
shortest
            | Bool -> (a -> Bool) -> Maybe a -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= Map Int a
dist Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
u) Maybe a
shortest = Map Int a -> [Int] -> Maybe a -> (Map Int a, Maybe a)
walk Map Int a
dist [Int]
us Maybe a
shortest
            | Bool
otherwise = Map Int a -> [Int] -> Maybe a -> [Int] -> (Map Int a, Maybe a)
foldNeighbors Map Int a
dist [Int]
us Maybe a
shortest (Int -> [Int]
adj Int
u)
          where
            foldNeighbors :: Map Int a -> [Int] -> Maybe a -> [Int] -> (Map Int a, Maybe a)
foldNeighbors Map Int a
d [Int]
queue Maybe a
found [] = Map Int a -> [Int] -> Maybe a -> (Map Int a, Maybe a)
walk Map Int a
d [Int]
queue Maybe a
found
            foldNeighbors Map Int a
d [Int]
queue Maybe a
found (Int
v : [Int]
vs) = case Int -> Map Int Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
v Map Int Int
reverseMatching of
                Maybe Int
Nothing -> Map Int a -> [Int] -> Maybe a -> [Int] -> (Map Int a, Maybe a)
foldNeighbors Map Int a
d [Int]
queue (a -> Maybe a
forall a. a -> Maybe a
Just (a -> (a -> a) -> Maybe a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Map Int a
d Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
u a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (a -> a -> a
forall a. Ord a => a -> a -> a
min (Map Int a
d Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
u a -> a -> a
forall a. Num a => a -> a -> a
+ a
1)) Maybe a
found)) [Int]
vs
                Just Int
partner
                    | Int -> Map Int a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.notMember Int
partner Map Int a
d ->
                        Map Int a -> [Int] -> Maybe a -> [Int] -> (Map Int a, Maybe a)
foldNeighbors (Int -> a -> Map Int a -> Map Int a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
partner (Map Int a
d Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
u a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) Map Int a
d) ([Int]
queue [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
partner]) Maybe a
found [Int]
vs
                    | Bool
otherwise -> Map Int a -> [Int] -> Maybe a -> [Int] -> (Map Int a, Maybe a)
foldNeighbors Map Int a
d [Int]
queue Maybe a
found [Int]
vs
    augmentPath :: Map Int a
-> a
-> Int
-> Map Int Int
-> Map Int Int
-> Maybe (Map Int Int, Map Int Int)
augmentPath Map Int a
dist a
shortest Int
u Map Int Int
matching Map Int Int
reverseMatching = [Int] -> Maybe (Map Int Int, Map Int Int)
try (Int -> [Int]
adj Int
u)
      where
        try :: [Int] -> Maybe (Map Int Int, Map Int Int)
try [] = Maybe (Map Int Int, Map Int Int)
forall a. Maybe a
Nothing
        try (Int
v : [Int]
vs) = case Int -> Map Int Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
v Map Int Int
reverseMatching of
            Maybe Int
Nothing
                | Int -> Map Int a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
u Map Int a
dist Maybe a -> Maybe a -> Bool
forall a. Eq a => a -> a -> Bool
== a -> Maybe a
forall a. a -> Maybe a
Just (a
shortest a -> a -> a
forall a. Num a => a -> a -> a
- a
1) ->
                    (Map Int Int, Map Int Int) -> Maybe (Map Int Int, Map Int Int)
forall a. a -> Maybe a
Just (Int -> Int -> Map Int Int -> Map Int Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
u Int
v Map Int Int
matching, Int -> Int -> Map Int Int -> Map Int Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
v Int
u Map Int Int
reverseMatching)
                | Bool
otherwise -> [Int] -> Maybe (Map Int Int, Map Int Int)
try [Int]
vs
            Just Int
partner | Int -> Map Int a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
partner Map Int a
dist Maybe a -> Maybe a -> Bool
forall a. Eq a => a -> a -> Bool
== ((a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (a -> a) -> Maybe a -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Map Int a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
u Map Int a
dist) -> case Map Int a
-> a
-> Int
-> Map Int Int
-> Map Int Int
-> Maybe (Map Int Int, Map Int Int)
augmentPath Map Int a
dist a
shortest Int
partner Map Int Int
matching Map Int Int
reverseMatching of
                Just (Map Int Int
m, Map Int Int
r) -> (Map Int Int, Map Int Int) -> Maybe (Map Int Int, Map Int Int)
forall a. a -> Maybe a
Just (Int -> Int -> Map Int Int -> Map Int Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
u Int
v Map Int Int
m, Int -> Int -> Map Int Int -> Map Int Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
v Int
u Map Int Int
r)
                Maybe (Map Int Int, Map Int Int)
Nothing -> [Int] -> Maybe (Map Int Int, Map Int Int)
try [Int]
vs
            Maybe Int
_ -> [Int] -> Maybe (Map Int Int, Map Int Int)
try [Int]
vs
    go :: Map Int Int -> Map Int Int -> Map Int Int
go Map Int Int
matching Map Int Int
reverseMatching = case Map Int Int -> Map Int Int -> (Map Int Int, Maybe Int)
forall {a}. Map Int a -> Map Int Int -> (Map Int Int, Maybe Int)
bfsLayer Map Int Int
matching Map Int Int
reverseMatching of
        (Map Int Int
_, Maybe Int
Nothing) -> Map Int Int
matching
        (Map Int Int
dist, Just Int
shortest) ->
            let (Map Int Int
next, Map Int Int
rev) =
                    ((Map Int Int, Map Int Int) -> Int -> (Map Int Int, Map Int Int))
-> (Map Int Int, Map Int Int)
-> [Int]
-> (Map Int Int, Map Int Int)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl'
                        (\(Map Int Int
m, Map Int Int
r) Int
u -> if Int -> Map Int Int -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member Int
u Map Int Int
m then (Map Int Int
m, Map Int Int
r) else (Map Int Int, Map Int Int)
-> ((Map Int Int, Map Int Int) -> (Map Int Int, Map Int Int))
-> Maybe (Map Int Int, Map Int Int)
-> (Map Int Int, Map Int Int)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Map Int Int
m, Map Int Int
r) (Map Int Int, Map Int Int) -> (Map Int Int, Map Int Int)
forall a. a -> a
id (Map Int Int
-> Int
-> Int
-> Map Int Int
-> Map Int Int
-> Maybe (Map Int Int, Map Int Int)
forall {a}.
(Num a, Eq a) =>
Map Int a
-> a
-> Int
-> Map Int Int
-> Map Int Int
-> Maybe (Map Int Int, Map Int Int)
augmentPath Map Int Int
dist Int
shortest Int
u Map Int Int
m Map Int Int
r))
                        (Map Int Int
matching, Map Int Int
reverseMatching)
                        [Int
0 .. Int
left Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
             in if Map Int Int -> Int
forall k a. Map k a -> Int
M.size Map Int Int
next Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Map Int Int -> Int
forall k a. Map k a -> Int
M.size Map Int Int
matching then Map Int Int
matching else Map Int Int -> Map Int Int -> Map Int Int
go Map Int Int
next Map Int Int
rev

{- | Rectangular Hungarian primal-dual assignment, rows <= columns <= 256.
Every row is assigned to a distinct column; exact integer costs may be negative.
-}
hungarian :: [[Integer]] -> Either MathError (Integer, [(Int, Int)])
hungarian :: [[Integer]] -> Either MathError (Integer, [(Int, Int)])
hungarian [[Integer]]
costs
    | Int
rows Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
cols Bool -> Bool -> Bool
|| Int
cols Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
256 Bool -> Bool -> Bool
|| ([Integer] -> Bool) -> [[Integer]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
cols) (Int -> Bool) -> ([Integer] -> Int) -> [Integer] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Integer]]
costs =
        MathError -> Either MathError (Integer, [(Int, Int)])
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"Hungarian matrix must be rectangular, rows <= columns <= 256")
    | Bool
otherwise =
        let (Residual
_, Residual
_, Map Int Int
assignment) = ((Residual, Residual, Map Int Int)
 -> Int -> (Residual, Residual, Map Int Int))
-> (Residual, Residual, Map Int Int)
-> [Int]
-> (Residual, Residual, Map Int Int)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (Residual, Residual, Map Int Int)
-> Int -> (Residual, Residual, Map Int Int)
addRow (Residual
forall k a. Map k a
M.empty, Residual
forall k a. Map k a
M.empty, Map Int Int
forall k a. Map k a
M.empty) [Int
1 .. Int
rows]
            pairs :: [(Int, Int)]
pairs = [(Int, Int)] -> [(Int, Int)]
forall a. Ord a => [a] -> [a]
sort [(Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) | (Int
j, Int
i) <- Map Int Int -> [(Int, Int)]
forall k a. Map k a -> [(k, a)]
M.toList Map Int Int
assignment, Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0]
         in (Integer, [(Int, Int)]) -> Either MathError (Integer, [(Int, Int)])
forall a b. b -> Either a b
Right ([Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[[Integer]]
costs [[Integer]] -> Int -> [Integer]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [Integer] -> Int -> Integer
forall a. HasCallStack => [a] -> Int -> a
!! Int
j | (Int
i, Int
j) <- [(Int, Int)]
pairs], [(Int, Int)]
pairs)
  where
    rows :: Int
rows = [[Integer]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Integer]]
costs
    cols :: Int
cols = if [[Integer]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Integer]]
costs then Int
0 else [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[Integer]]
costs [[Integer]] -> Int -> [Integer]
forall a. HasCallStack => [a] -> Int -> a
!! Int
0)
    value :: k -> Map k a -> a
value k
key = a -> k -> Map k a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 k
key
    addRow :: (Residual, Residual, Map Int Int)
-> Int -> (Residual, Residual, Map Int Int)
addRow (Residual
u, Residual
v, Map Int Int
p) Int
row = Int
-> Residual
-> Residual
-> Map Int Int
-> Set Int
-> Residual
-> Map Int Int
-> (Residual, Residual, Map Int Int)
search Int
0 Residual
u Residual
v (Int -> Int -> Map Int Int -> Map Int Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
0 Int
row Map Int Int
p) Set Int
forall a. Set a
S.empty Residual
forall k a. Map k a
M.empty Map Int Int
forall k a. Map k a
M.empty
    search :: Int
-> Residual
-> Residual
-> Map Int Int
-> Set Int
-> Residual
-> Map Int Int
-> (Residual, Residual, Map Int Int)
search Int
j0 Residual
u Residual
v Map Int Int
p Set Int
used Residual
slack Map Int Int
way =
        let used' :: Set Int
used' = Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
S.insert Int
j0 Set Int
used
            i0 :: Int
i0 = Map Int Int
p Map Int Int -> Int -> Int
forall k a. Ord k => Map k a -> k -> a
M.! Int
j0
            candidates :: [Int]
candidates = [Int
j | Int
j <- [Int
1 .. Int
cols], Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
S.notMember Int
j Set Int
used']
            (Residual
slack', Map Int Int
way') =
                ((Residual, Map Int Int) -> Int -> (Residual, Map Int Int))
-> (Residual, Map Int Int) -> [Int] -> (Residual, Map Int Int)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl'
                    ( \(Residual
sl, Map Int Int
wa) Int
j ->
                        let reduced :: Integer
reduced = [[Integer]]
costs [[Integer]] -> Int -> [Integer]
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Integer] -> Int -> Integer
forall a. HasCallStack => [a] -> Int -> a
!! (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Residual -> Integer
forall {k} {a}. (Ord k, Num a) => k -> Map k a -> a
value Int
i0 Residual
u Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Residual -> Integer
forall {k} {a}. (Ord k, Num a) => k -> Map k a -> a
value Int
j Residual
v
                         in if Bool -> (Integer -> Bool) -> Maybe Integer -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Integer
reduced Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<) (Int -> Residual -> Maybe Integer
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
j Residual
sl)
                                then (Int -> Integer -> Residual -> Residual
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
j Integer
reduced Residual
sl, Int -> Int -> Map Int Int -> Map Int Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
j Int
j0 Map Int Int
wa)
                                else (Residual
sl, Map Int Int
wa)
                    )
                    (Residual
slack, Map Int Int
way)
                    [Int]
candidates
            (Integer
delta, Int
j1) = [(Integer, Int)] -> (Integer, Int)
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [(Residual
slack' Residual -> Int -> Integer
forall k a. Ord k => Map k a -> k -> a
M.! Int
j, Int
j) | Int
j <- [Int]
candidates]
            u' :: Residual
u' = (Residual -> Int -> Residual) -> Residual -> [Int] -> Residual
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\Residual
m Int
j -> (Integer -> Integer -> Integer)
-> Int -> Integer -> Residual -> Residual
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) (Map Int Int
p Map Int Int -> Int -> Int
forall k a. Ord k => Map k a -> k -> a
M.! Int
j) Integer
delta Residual
m) Residual
u (Set Int -> [Int]
forall a. Set a -> [a]
S.toList Set Int
used')
            v' :: Residual
v' = (Residual -> Int -> Residual) -> Residual -> [Int] -> Residual
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\Residual
m Int
j -> (Integer -> Integer -> Integer)
-> Int -> Integer -> Residual -> Residual
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) Int
j (-Integer
delta) Residual
m) Residual
v (Set Int -> [Int]
forall a. Set a -> [a]
S.toList Set Int
used')
            sl' :: Residual
sl' = (Int -> Integer -> Integer) -> Residual -> Residual
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey (\Int
j Integer
d -> if Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Int
j Set Int
used' then Integer
d else Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
delta) Residual
slack'
            augmentAssignment :: Int -> Map Int a -> Map Int a
augmentAssignment Int
j Map Int a
assignment
                | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Map Int a
assignment
                | Bool
otherwise =
                    let previous :: Int
previous = Map Int Int
way' Map Int Int -> Int -> Int
forall k a. Ord k => Map k a -> k -> a
M.! Int
j
                     in Int -> Map Int a -> Map Int a
augmentAssignment Int
previous (Int -> a -> Map Int a -> Map Int a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
j (Map Int a
assignment Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
previous) Map Int a
assignment)
         in if Int -> Map Int Int -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.notMember Int
j1 Map Int Int
p
                then (Residual
u', Residual
v', Int -> Map Int Int -> Map Int Int
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Int
0 (Int -> Map Int Int -> Map Int Int
forall {a}. Int -> Map Int a -> Map Int a
augmentAssignment Int
j1 Map Int Int
p))
                else Int
-> Residual
-> Residual
-> Map Int Int
-> Set Int
-> Residual
-> Map Int Int
-> (Residual, Residual, Map Int Int)
search Int
j1 Residual
u' Residual
v' Map Int Int
p Set Int
used' Residual
sl' Map Int Int
way'