{-# LANGUAGE ScopedTypeVariables #-}
{-|
== Dinic's algorithm, or Dinitz's algorithm

An algorithm to find the maximum flow in a flow network.

Sources:

* Y. Dinitz, "Dinitz' Algorithm: The Original Version and Even's Version", 2006
  https://www.cs.bgu.ac.il/~dinitz/Papers/Dinitz_alg.pdf
* AC Library
  https://github.com/atcoder/ac-library/blob/master/atcoder/maxflow.hpp

-}

{-
Implementation notes:
* This implementation is close to B. Cherkassky's implementation in Dinitz's paper above, and the
  implementation in AC Library.
* The flow array contains flows of edges in the residual graph, forward and backward edges occupy
  even and odd positions.
-}

module Dinic
    ( FlowEdge(..)
    , FlowResult(..)
    , ToEdge(..)
    , Flow
    , EdgeIndex
    , dinic
    ) where

import Control.DeepSeq
import Control.Monad.ST
import Data.Array.ST
import Data.Array.Unboxed
import Data.Array.Unsafe
import Data.Bits
import Data.Graph

import Misc ( modifyArray )

type Flow = Int
data FlowEdge = FlowEdge { FlowEdge -> Vertex
from_ :: !Vertex, FlowEdge -> Vertex
to_ :: !Vertex, FlowEdge -> Vertex
cap_ :: !Flow } deriving (FlowEdge -> FlowEdge -> Bool
(FlowEdge -> FlowEdge -> Bool)
-> (FlowEdge -> FlowEdge -> Bool) -> Eq FlowEdge
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: FlowEdge -> FlowEdge -> Bool
$c/= :: FlowEdge -> FlowEdge -> Bool
== :: FlowEdge -> FlowEdge -> Bool
$c== :: FlowEdge -> FlowEdge -> Bool
Eq, Vertex -> FlowEdge -> ShowS
[FlowEdge] -> ShowS
FlowEdge -> String
(Vertex -> FlowEdge -> ShowS)
-> (FlowEdge -> String) -> ([FlowEdge] -> ShowS) -> Show FlowEdge
forall a.
(Vertex -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [FlowEdge] -> ShowS
$cshowList :: [FlowEdge] -> ShowS
show :: FlowEdge -> String
$cshow :: FlowEdge -> String
showsPrec :: Vertex -> FlowEdge -> ShowS
$cshowsPrec :: Vertex -> FlowEdge -> ShowS
Show)

type EdgeIndex = Int
data ToEdge = ToEdge { ToEdge -> Vertex
to__ :: !Vertex, ToEdge -> Vertex
edgeIndex_ :: !EdgeIndex }

data FlowResult = FlowResult { FlowResult -> Vertex
getFlow :: !Flow , FlowResult -> [Vertex]
getFlows :: [Flow] , FlowResult -> [Bool]
getMinCut :: [Bool] }

-- | Runs Dinic's algorithm on the graph made up of the given FlowEdges. Returns a FlowResult, describing
-- a max flow configuration with
-- 1. the max flow value
-- 2. flow values of the edges in the order in which they were given
-- 3. a list of Bools for each edge indicating whether the edge is in a min-cut
-- O(V^2 E) in the general case.
-- O(min(V^(2/3), E^(1/2)) E) for unit capacity graphs.
-- O(V^(1/2) E) for unit networks, such as in maximum bipartite matching.
dinic :: Bounds -> [FlowEdge] -> Vertex -> Vertex -> FlowResult
dinic :: Bounds -> [FlowEdge] -> Vertex -> Vertex -> FlowResult
dinic Bounds
_ [FlowEdge]
_ Vertex
src Vertex
sink | Vertex
src Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
sink = String -> FlowResult
forall a. HasCallStack => String -> a
error String
"src == sink"
dinic Bounds
bnds [FlowEdge]
es Vertex
src Vertex
sink = Vertex -> [Vertex] -> [Bool] -> FlowResult
FlowResult Vertex
maxFlow [Vertex]
flows [Bool]
minCut where
    m :: Vertex
m = [FlowEdge] -> Vertex
forall (t :: * -> *) a. Foldable t => t a -> Vertex
length [FlowEdge]
es
    Array Vertex [ToEdge]
g :: Array Vertex [ToEdge] = ([ToEdge] -> ToEdge -> [ToEdge])
-> [ToEdge]
-> Bounds
-> [(Vertex, ToEdge)]
-> Array Vertex [ToEdge]
forall (a :: * -> * -> *) e i e'.
(IArray a e, Ix i) =>
(e -> e' -> e) -> e -> (i, i) -> [(i, e')] -> a i e
accumArray ((ToEdge -> [ToEdge] -> [ToEdge]) -> [ToEdge] -> ToEdge -> [ToEdge]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) [] Bounds
bnds ([(Vertex, ToEdge)] -> Array Vertex [ToEdge])
-> [(Vertex, ToEdge)] -> Array Vertex [ToEdge]
forall a b. (a -> b) -> a -> b
$ do
        (Vertex
i, FlowEdge Vertex
u Vertex
v Vertex
_) <- [Vertex] -> [FlowEdge] -> [(Vertex, FlowEdge)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Vertex
0..] [FlowEdge]
es
        [(Vertex
u, Vertex -> Vertex -> ToEdge
ToEdge Vertex
v (Vertex
2Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
*Vertex
i)), (Vertex
v, Vertex -> Vertex -> ToEdge
ToEdge Vertex
u (Vertex
2Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
*Vertex
iVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+Vertex
1))]
    flows :: [Vertex]
flows = (Vertex -> Vertex) -> [Vertex] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map (UArray Vertex Vertex
flowaUArray Vertex Vertex -> Vertex -> Vertex
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!) [Vertex
0,Vertex
2..Vertex
2Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
*Vertex
mVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
-Vertex
2]
    minCut :: [Bool]
minCut = [UArray Vertex Vertex
lvlUArray Vertex Vertex -> Vertex -> Vertex
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
u Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== -Vertex
1 Bool -> Bool -> Bool
&& UArray Vertex Vertex
lvlUArray Vertex Vertex -> Vertex -> Vertex
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
/= -Vertex
1 | FlowEdge Vertex
u Vertex
v Vertex
_ <- [FlowEdge]
es] where
        lvl :: UArray Vertex Vertex
lvl = (forall s. ST s (UArray Vertex Vertex)) -> UArray Vertex Vertex
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (UArray Vertex Vertex)) -> UArray Vertex Vertex)
-> (forall s. ST s (UArray Vertex Vertex)) -> UArray Vertex Vertex
forall a b. (a -> b) -> a -> b
$ (Vertex -> ST s Bool) -> ST s (UArray Vertex Vertex)
forall s. (Vertex -> ST s Bool) -> ST s (UArray Vertex Vertex)
dinicLevels (Bool -> ST s Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> ST s Bool) -> (Vertex -> Bool) -> Vertex -> ST s Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
==Vertex
0) (Vertex -> Bool) -> (Vertex -> Vertex) -> Vertex -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UArray Vertex Vertex
flowaUArray Vertex Vertex -> Vertex -> Vertex
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!))

    (Vertex
maxFlow, UArray Vertex Vertex
flowa :: UArray Int Flow) = (forall s. ST s (Vertex, UArray Vertex Vertex))
-> (Vertex, UArray Vertex Vertex)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vertex, UArray Vertex Vertex))
 -> (Vertex, UArray Vertex Vertex))
-> (forall s. ST s (Vertex, UArray Vertex Vertex))
-> (Vertex, UArray Vertex Vertex)
forall a b. (a -> b) -> a -> b
$ do
        STUArray s Vertex Vertex
flow :: STUArray s EdgeIndex Flow <-
            Bounds -> [Vertex] -> ST s (STUArray s Vertex Vertex)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> [e] -> m (a i e)
newListArray (Vertex
0, Vertex
2Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
*Vertex
mVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
-Vertex
1) ([Vertex] -> ST s (STUArray s Vertex Vertex))
-> [Vertex] -> ST s (STUArray s Vertex Vertex)
forall a b. (a -> b) -> a -> b
$ [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Vertex
0, Vertex
c] | FlowEdge Vertex
_ Vertex
_ Vertex
c <- [FlowEdge]
es]
        let runDinic :: ST s Vertex
runDinic = do
                UArray Vertex Vertex
lvl <- (Vertex -> ST s Bool) -> ST s (UArray Vertex Vertex)
forall s. (Vertex -> ST s Bool) -> ST s (UArray Vertex Vertex)
dinicLevels ((Vertex -> Bool) -> ST s Vertex -> ST s Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
==Vertex
0) (ST s Vertex -> ST s Bool)
-> (Vertex -> ST s Vertex) -> Vertex -> ST s Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Vertex Vertex
flow)
                Vertex
f <- (Vertex -> Vertex -> Bool)
-> STUArray s Vertex Vertex -> ST s Vertex
forall s.
(Vertex -> Vertex -> Bool)
-> STUArray s Vertex Vertex -> ST s Vertex
dinicAugment (\Vertex
u Vertex
v -> UArray Vertex Vertex
lvlUArray Vertex Vertex -> Vertex -> Vertex
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== UArray Vertex Vertex
lvlUArray Vertex Vertex -> Vertex -> Vertex
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
u Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
1) STUArray s Vertex Vertex
flow
                if Vertex
f Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
0 then Vertex -> ST s Vertex
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
0 else (Vertex
fVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+) (Vertex -> Vertex) -> ST s Vertex -> ST s Vertex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ST s Vertex
runDinic
        (,) (Vertex -> UArray Vertex Vertex -> (Vertex, UArray Vertex Vertex))
-> ST s Vertex
-> ST s (UArray Vertex Vertex -> (Vertex, UArray Vertex Vertex))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ST s Vertex
runDinic ST s (UArray Vertex Vertex -> (Vertex, UArray Vertex Vertex))
-> ST s (UArray Vertex Vertex)
-> ST s (Vertex, UArray Vertex Vertex)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> STUArray s Vertex Vertex -> ST s (UArray Vertex Vertex)
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze STUArray s Vertex Vertex
flow

    dinicLevels :: forall s. (EdgeIndex -> ST s Bool) -> ST s (UArray Vertex Int)
    dinicLevels :: (Vertex -> ST s Bool) -> ST s (UArray Vertex Vertex)
dinicLevels Vertex -> ST s Bool
sat = do
        STUArray s Vertex Vertex
lvl :: STUArray s Vertex Int <- Bounds -> Vertex -> ST s (STUArray s Vertex Vertex)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Array Vertex [ToEdge] -> Bounds
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds Array Vertex [ToEdge]
g) (-Vertex
1)
        let visit :: Vertex -> [Vertex] -> ST s ()
visit Vertex
_ [] = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            visit Vertex
d [Vertex]
q = [ToEdge] -> [Vertex] -> [Vertex] -> ST s [Vertex]
go [] [Vertex]
q [] ST s [Vertex] -> ([Vertex] -> ST s ()) -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Vertex -> [Vertex] -> ST s ()
visit (Vertex
d Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1) where
                go :: [ToEdge] -> [Vertex] -> [Vertex] -> ST s [Vertex]
go [] []     [Vertex]
acc = [Vertex] -> ST s [Vertex]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Vertex]
acc
                go [] (Vertex
u:[Vertex]
us) [Vertex]
acc = [ToEdge] -> [Vertex] -> [Vertex] -> ST s [Vertex]
go (Array Vertex [ToEdge]
gArray Vertex [ToEdge] -> Vertex -> [ToEdge]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
u) [Vertex]
us [Vertex]
acc
                go (ToEdge Vertex
v Vertex
i:[ToEdge]
ts) [Vertex]
us [Vertex]
acc = do
                    Vertex
l <- STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Vertex Vertex
lvl Vertex
v
                    Bool
s <- Vertex -> ST s Bool
sat Vertex
i
                    if Vertex
l Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
/= -Vertex
1 Bool -> Bool -> Bool
|| Bool
s
                        then [ToEdge] -> [Vertex] -> [Vertex] -> ST s [Vertex]
go [ToEdge]
ts [Vertex]
us [Vertex]
acc
                        else do
                            STUArray s Vertex Vertex -> Vertex -> Vertex -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Vertex Vertex
lvl Vertex
v Vertex
d
                            if Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
src then [Vertex] -> ST s [Vertex]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [] else [ToEdge] -> [Vertex] -> [Vertex] -> ST s [Vertex]
go [ToEdge]
ts [Vertex]
us (Vertex
vVertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
:[Vertex]
acc)
        STUArray s Vertex Vertex -> Vertex -> Vertex -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Vertex Vertex
lvl Vertex
sink Vertex
0
        Vertex -> [Vertex] -> ST s ()
visit Vertex
1 [Vertex
sink]
        STUArray s Vertex Vertex -> ST s (UArray Vertex Vertex)
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze STUArray s Vertex Vertex
lvl
    {-# INLINE dinicLevels #-}

    dinicAugment :: forall s. (Vertex -> Vertex -> Bool) -> STUArray s EdgeIndex Flow -> ST s Flow
    dinicAugment :: (Vertex -> Vertex -> Bool)
-> STUArray s Vertex Vertex -> ST s Vertex
dinicAugment Vertex -> Vertex -> Bool
nxtLvl STUArray s Vertex Vertex
flow = do
        STArray s Vertex [ToEdge]
g' :: STArray s Vertex [ToEdge] <- Array Vertex [ToEdge] -> ST s (STArray s Vertex [ToEdge])
forall i (a :: * -> * -> *) e (b :: * -> * -> *) (m :: * -> *).
(Ix i, IArray a e, MArray b e m) =>
a i e -> m (b i e)
thaw Array Vertex [ToEdge]
g
        let go :: Vertex -> Vertex -> ST s Vertex
go Vertex
_ Vertex
0 = Vertex -> ST s Vertex
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
0
            go Vertex
u Vertex
fup | Vertex
u Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
sink = Vertex -> ST s Vertex
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
fup
            go Vertex
u Vertex
fup = STArray s Vertex [ToEdge] -> Vertex -> ST s [ToEdge]
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Vertex [ToEdge]
g' Vertex
u ST s [ToEdge] -> ([ToEdge] -> ST s Vertex) -> ST s Vertex
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [ToEdge] -> ST s Vertex
go' where
                go' :: [ToEdge] -> ST s Vertex
go' [] = Vertex -> ST s Vertex
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
0 :: ST s Flow
                go' (ToEdge Vertex
v Vertex
_:[ToEdge]
ts) | Bool -> Bool
not (Vertex -> Vertex -> Bool
nxtLvl Vertex
u Vertex
v) = STArray s Vertex [ToEdge] -> Vertex -> [ToEdge] -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Vertex [ToEdge]
g' Vertex
u [ToEdge]
ts ST s () -> ST s Vertex -> ST s Vertex
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Vertex -> Vertex -> ST s Vertex
go Vertex
u Vertex
fup
                go' (ToEdge Vertex
v Vertex
i:[ToEdge]
ts) = do
                    Vertex
f <- STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Vertex Vertex
flow (Vertex -> Vertex -> Vertex
forall a. Bits a => a -> a -> a
xor Vertex
i Vertex
1)
                    Vertex
fdn <- Vertex -> Vertex -> ST s Vertex
go Vertex
v (Vertex -> Vertex -> Vertex
forall a. Ord a => a -> a -> a
min Vertex
fup Vertex
f)
                    STUArray s Vertex Vertex -> Vertex -> (Vertex -> Vertex) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray STUArray s Vertex Vertex
flow Vertex
i (Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+Vertex
fdn) 
                    STUArray s Vertex Vertex -> Vertex -> (Vertex -> Vertex) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray STUArray s Vertex Vertex
flow (Vertex -> Vertex -> Vertex
forall a. Bits a => a -> a -> a
xor Vertex
i Vertex
1) (Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+(-Vertex
fdn))
                    if Vertex
fdn Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
fup
                        then Vertex -> ST s Vertex
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
fdn
                        else STArray s Vertex [ToEdge] -> Vertex -> [ToEdge] -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Vertex [ToEdge]
g' Vertex
u [ToEdge]
ts ST s () -> ST s Vertex -> ST s Vertex
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Vertex
fdnVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+) (Vertex -> Vertex) -> ST s Vertex -> ST s Vertex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vertex -> Vertex -> ST s Vertex
go Vertex
u (Vertex
fup Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
fdn)
        Vertex -> Vertex -> ST s Vertex
go Vertex
src Vertex
forall a. Bounded a => a
maxBound
    {-# INLINE dinicAugment #-}

--------------------------------------------------------------------------------
-- For tests

instance NFData FlowEdge where
    rnf :: FlowEdge -> ()
rnf = FlowEdge -> ()
forall a. a -> ()
rwhnf