{-# LANGUAGE RecordWildCards #-}
module AtCoder.MinCostFlow
(
McfGraph (nG),
new,
addEdge,
addEdge_,
flow,
maxFlow,
slope,
getEdge,
edges,
unsafeEdges,
)
where
import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.Internal.Buffer qualified as ACIB
import AtCoder.Internal.GrowVec qualified as ACIGV
import AtCoder.Internal.McfCsr qualified as ACIMCSR
import AtCoder.Internal.MinHeap qualified as ACIMH
import Control.Monad (unless, when)
import Control.Monad.Fix (fix)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Bit (Bit (..))
import Data.Maybe (fromJust)
import Data.Primitive.MutVar (readMutVar)
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Base qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Stack (HasCallStack)
data McfGraph s cap cost = McfGraph
{
forall s cap cost. McfGraph s cap cost -> Int
nG :: {-# UNPACK #-} !Int,
forall s cap cost.
McfGraph s cap cost -> GrowVec s (Int, Int, cap, cap, cost)
edgesG :: !(ACIGV.GrowVec s (Int, Int, cap, cap, cost))
}
{-# INLINE new #-}
new :: (PrimMonad m, VU.Unbox cap, VU.Unbox cost) => Int -> m (McfGraph (PrimState m) cap cost)
new :: forall (m :: * -> *) cap cost.
(PrimMonad m, Unbox cap, Unbox cost) =>
Int -> m (McfGraph (PrimState m) cap cost)
new Int
nG = do
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
edgesG <- Int -> m (GrowVec (PrimState m) (Int, Int, cap, cap, cost))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (GrowVec (PrimState m) a)
ACIGV.new Int
0
McfGraph (PrimState m) cap cost
-> m (McfGraph (PrimState m) cap cost)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure McfGraph {Int
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
nG :: Int
edgesG :: GrowVec (PrimState m) (Int, Int, cap, cap, cost)
nG :: Int
edgesG :: GrowVec (PrimState m) (Int, Int, cap, cap, cost)
..}
{-# INLINE addEdge #-}
addEdge ::
(HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
McfGraph (PrimState m) cap cost ->
Int ->
Int ->
cap ->
cost ->
m Int
addEdge :: forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Num cap, Ord cap, Unbox cap, Num cost,
Ord cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> cost -> m Int
addEdge McfGraph {Int
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
nG :: forall s cap cost. McfGraph s cap cost -> Int
edgesG :: forall s cap cost.
McfGraph s cap cost -> GrowVec s (Int, Int, cap, cap, cost)
nG :: Int
edgesG :: GrowVec (PrimState m) (Int, Int, cap, cap, cost)
..} Int
from Int
to cap
cap cost
cost = do
let !()
_ = HasCallStack => String -> String -> Int -> String -> Int -> ()
String -> String -> Int -> String -> Int -> ()
ACIA.checkCustom String
"AtCoder.MinCostFlow.addEdge" String
"`from` vertex" Int
from String
"the number of vertices" Int
nG
let !()
_ = HasCallStack => String -> String -> Int -> String -> Int -> ()
String -> String -> Int -> String -> Int -> ()
ACIA.checkCustom String
"AtCoder.MinCostFlow.addEdge" String
"`to` vertex" Int
to String
"the number of vertices" Int
nG
let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (cap
0 cap -> cap -> Bool
forall a. Ord a => a -> a -> Bool
<= cap
cap) String
"AtCoder.MinCostFlow.addEdge: given invalid edge `cap` less than `0`"
let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (cost
0 cost -> cost -> Bool
forall a. Ord a => a -> a -> Bool
<= cost
cost) String
"AtCoder.MinCostFlow.addEdge: given invalid edge `cost` less than `0`"
Int
m <- GrowVec (PrimState m) (Int, Int, cap, cap, cost) -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m Int
ACIGV.length GrowVec (PrimState m) (Int, Int, cap, cap, cost)
edgesG
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
-> (Int, Int, cap, cap, cost) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
ACIGV.pushBack GrowVec (PrimState m) (Int, Int, cap, cap, cost)
edgesG (Int
from, Int
to, cap
cap, cap
0, cost
cost)
Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
m
{-# INLINE addEdge_ #-}
addEdge_ ::
(HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
McfGraph (PrimState m) cap cost ->
Int ->
Int ->
cap ->
cost ->
m ()
addEdge_ :: forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Num cap, Ord cap, Unbox cap, Num cost,
Ord cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> cost -> m ()
addEdge_ McfGraph (PrimState m) cap cost
graph Int
from Int
to cap
cap cost
cost = do
Int
_ <- McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> cost -> m Int
forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Num cap, Ord cap, Unbox cap, Num cost,
Ord cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> cost -> m Int
addEdge McfGraph (PrimState m) cap cost
graph Int
from Int
to cap
cap cost
cost
() -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE flow #-}
flow ::
(HasCallStack, PrimMonad m, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>
McfGraph (PrimState m) cap cost ->
Int ->
Int ->
cap ->
m (cap, cost)
flow :: forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Integral cap, Ord cap, Unbox cap,
Num cost, Ord cost, Bounded cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> m (cap, cost)
flow McfGraph (PrimState m) cap cost
graph Int
s Int
t cap
flowLimit = do
Vector (cap, cost)
res <- McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> m (Vector (cap, cost))
forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Integral cap, Ord cap, Unbox cap,
Num cost, Ord cost, Bounded cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> m (Vector (cap, cost))
slope McfGraph (PrimState m) cap cost
graph Int
s Int
t cap
flowLimit
(cap, cost) -> m (cap, cost)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((cap, cost) -> m (cap, cost)) -> (cap, cost) -> m (cap, cost)
forall a b. (a -> b) -> a -> b
$ Vector (cap, cost) -> (cap, cost)
forall (v :: * -> *) a. Vector v a => v a -> a
VG.last Vector (cap, cost)
res
{-# INLINE maxFlow #-}
maxFlow ::
(HasCallStack, PrimMonad m, Integral cap, Ord cap, Bounded cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>
McfGraph (PrimState m) cap cost ->
Int ->
Int ->
m (cap, cost)
maxFlow :: forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Integral cap, Ord cap, Bounded cap,
Unbox cap, Num cost, Ord cost, Bounded cost, Unbox cost) =>
McfGraph (PrimState m) cap cost -> Int -> Int -> m (cap, cost)
maxFlow McfGraph (PrimState m) cap cost
graph Int
s Int
t = do
Vector (cap, cost)
res <- McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> m (Vector (cap, cost))
forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Integral cap, Ord cap, Unbox cap,
Num cost, Ord cost, Bounded cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> m (Vector (cap, cost))
slope McfGraph (PrimState m) cap cost
graph Int
s Int
t cap
forall a. Bounded a => a
maxBound
(cap, cost) -> m (cap, cost)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((cap, cost) -> m (cap, cost)) -> (cap, cost) -> m (cap, cost)
forall a b. (a -> b) -> a -> b
$ Vector (cap, cost) -> (cap, cost)
forall (v :: * -> *) a. Vector v a => v a -> a
VG.last Vector (cap, cost)
res
{-# INLINE slope #-}
slope ::
(HasCallStack, PrimMonad m, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>
McfGraph (PrimState m) cap cost ->
Int ->
Int ->
cap ->
m (VU.Vector (cap, cost))
slope :: forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Integral cap, Ord cap, Unbox cap,
Num cost, Ord cost, Bounded cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> Int -> Int -> cap -> m (Vector (cap, cost))
slope McfGraph {Int
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
nG :: forall s cap cost. McfGraph s cap cost -> Int
edgesG :: forall s cap cost.
McfGraph s cap cost -> GrowVec s (Int, Int, cap, cap, cost)
nG :: Int
edgesG :: GrowVec (PrimState m) (Int, Int, cap, cap, cost)
..} Int
s Int
t cap
flowLimit = do
let !()
_ = HasCallStack => String -> String -> Int -> String -> Int -> ()
String -> String -> Int -> String -> Int -> ()
ACIA.checkCustom String
"AtCoder.MinCostFlow.slope" String
"`source` vertex" Int
s String
"the number of vertices" Int
nG
let !()
_ = HasCallStack => String -> String -> Int -> String -> Int -> ()
String -> String -> Int -> String -> Int -> ()
ACIA.checkCustom String
"AtCoder.MinCostFlow.slope" String
"`sink` vertex" Int
t String
"the number of vertices" Int
nG
let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
t) String
"AtCoder.MinCostFlow.slope: `source` and `sink` vertex must be distict"
edges_ :: Vector (Int, Int, cap, cap, cost)
edges_@(VU.V_5 Int
_ Vector Int
_ Vector Int
_ Vector cap
caps Vector cap
_ Vector cost
_) <- GrowVec (PrimState m) (Int, Int, cap, cap, cost)
-> m (Vector (Int, Int, cap, cap, cost))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
ACIGV.unsafeFreeze GrowVec (PrimState m) (Int, Int, cap, cap, cost)
edgesG
(!Vector Int
edgeIdx, !Csr (PrimState m) cap cost
g) <- Int
-> Vector (Int, Int, cap, cap, cost)
-> m (Vector Int, Csr (PrimState m) cap cost)
forall cap cost (m :: * -> *).
(HasCallStack, Num cap, Unbox cap, Unbox cost, Num cost,
PrimMonad m) =>
Int
-> Vector (Int, Int, cap, cap, cost)
-> m (Vector Int, Csr (PrimState m) cap cost)
ACIMCSR.build Int
nG Vector (Int, Int, cap, cap, cost)
edges_
Vector (cap, cost)
result <- Csr (PrimState m) cap cost
-> Int -> Int -> Int -> cap -> m (Vector (cap, cost))
forall cap cost (m :: * -> *).
(HasCallStack, PrimMonad m, Integral cap, Ord cap, Unbox cap,
Num cost, Ord cost, Bounded cost, Unbox cost) =>
Csr (PrimState m) cap cost
-> Int -> Int -> Int -> cap -> m (Vector (cap, cost))
internalSlopeMCF Csr (PrimState m) cap cost
g Int
nG Int
s Int
t cap
flowLimit
(VUM.MV_5 Int
_ MVector (PrimState m) Int
_ MVector (PrimState m) Int
_ MVector (PrimState m) cap
_ MVector (PrimState m) cap
flows MVector (PrimState m) cost
_) <- MutVar
(PrimState m) (MVector (PrimState m) (Int, Int, cap, cap, cost))
-> m (MVector (PrimState m) (Int, Int, cap, cap, cost))
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar (MutVar
(PrimState m) (MVector (PrimState m) (Int, Int, cap, cap, cost))
-> m (MVector (PrimState m) (Int, Int, cap, cap, cost)))
-> MutVar
(PrimState m) (MVector (PrimState m) (Int, Int, cap, cap, cost))
-> m (MVector (PrimState m) (Int, Int, cap, cap, cost))
forall a b. (a -> b) -> a -> b
$ GrowVec (PrimState m) (Int, Int, cap, cap, cost)
-> MutVar
(PrimState m) (MVector (PrimState m) (Int, Int, cap, cap, cost))
forall s a. GrowVec s a -> MutVar s (MVector s a)
ACIGV.vecGV GrowVec (PrimState m) (Int, Int, cap, cap, cost)
edgesG
Vector (cap, Int) -> (Int -> (cap, Int) -> m ()) -> m ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ (Vector cap -> Vector Int -> Vector (cap, Int)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector cap
caps Vector Int
edgeIdx) ((Int -> (cap, Int) -> m ()) -> m ())
-> (Int -> (cap, Int) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
v (!cap
cap1, !Int
iEdge) -> do
cap
cap2 <- MVector (PrimState m) cap -> Int -> m cap
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read (Csr (PrimState m) cap cost -> MVector (PrimState m) cap
forall s cap cost. Csr s cap cost -> MVector s cap
ACIMCSR.capCsr Csr (PrimState m) cap cost
g) Int
iEdge
MVector (PrimState m) cap -> Int -> cap -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) cap
flows Int
v (cap -> m ()) -> cap -> m ()
forall a b. (a -> b) -> a -> b
$! cap
cap1 cap -> cap -> cap
forall a. Num a => a -> a -> a
- cap
cap2
Vector (cap, cost) -> m (Vector (cap, cost))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector (cap, cost)
result
{-# INLINE internalSlopeMCF #-}
internalSlopeMCF ::
forall cap cost m.
(HasCallStack, PrimMonad m, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>
ACIMCSR.Csr (PrimState m) cap cost ->
Int ->
Int ->
Int ->
cap ->
m (VU.Vector (cap, cost))
internalSlopeMCF :: forall cap cost (m :: * -> *).
(HasCallStack, PrimMonad m, Integral cap, Ord cap, Unbox cap,
Num cost, Ord cost, Bounded cost, Unbox cost) =>
Csr (PrimState m) cap cost
-> Int -> Int -> Int -> cap -> m (Vector (cap, cost))
internalSlopeMCF csr :: Csr (PrimState m) cap cost
csr@ACIMCSR.Csr {MVector (PrimState m) cap
Vector cost
Vector Int
capCsr :: forall s cap cost. Csr s cap cost -> MVector s cap
startCsr :: Vector Int
toCsr :: Vector Int
revCsr :: Vector Int
capCsr :: MVector (PrimState m) cap
costCsr :: Vector cost
startCsr :: forall s cap cost. Csr s cap cost -> Vector Int
toCsr :: forall s cap cost. Csr s cap cost -> Vector Int
revCsr :: forall s cap cost. Csr s cap cost -> Vector Int
costCsr :: forall s cap cost. Csr s cap cost -> Vector cost
..} Int
n Int
s Int
t cap
flowLimit = do
MVector (PrimState m) cost
duals <- Int -> cost -> m (MVector (PrimState m) cost)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
n cost
0
MVector (PrimState m) cost
dists <- Int -> m (MVector (PrimState m) cost)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n :: m (VUM.MVector (PrimState m) cost)
MVector (PrimState m) Int
prevE <- Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n :: m (VUM.MVector (PrimState m) Int)
MVector (PrimState m) Bit
vis <- Int -> m (MVector (PrimState m) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n :: m (VUM.MVector (PrimState m) Bit)
let nEdges :: Int
nEdges = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
toCsr
Buffer (PrimState m) Int
queMin <- Int -> m (Buffer (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Buffer (PrimState m) a)
ACIB.new Int
nEdges :: m (ACIB.Buffer (PrimState m) Int)
Heap (PrimState m) (cost, Int)
heap <- Int -> m (Heap (PrimState m) (cost, Int))
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Int -> m (Heap (PrimState m) a)
ACIMH.new Int
nEdges :: m (ACIMH.Heap (PrimState m) (cost, Int))
let dualRef :: m Bool
dualRef = do
MVector (PrimState m) cost -> cost -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
VGM.set MVector (PrimState m) cost
dists (cost -> m ()) -> cost -> m ()
forall a b. (a -> b) -> a -> b
$ forall a. Bounded a => a
maxBound @cost
MVector (PrimState m) Bit -> Bit -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
VGM.set MVector (PrimState m) Bit
vis (Bit -> m ()) -> Bit -> m ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False
Buffer (PrimState m) Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m ()
ACIB.clear Buffer (PrimState m) Int
queMin
Heap (PrimState m) (cost, Int) -> m ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Heap (PrimState m) a -> m ()
ACIMH.clear Heap (PrimState m) (cost, Int)
heap
MVector (PrimState m) cost -> Int -> cost -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) cost
dists Int
s cost
0
Buffer (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
ACIB.pushBack Buffer (PrimState m) Int
queMin Int
s
(m () -> m ()) -> m ()
forall a. (a -> a) -> a
fix ((m () -> m ()) -> m ()) -> (m () -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \m ()
loop -> do
Bool
b1 <- Buffer (PrimState m) Int -> m Bool
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m Bool
ACIB.null Buffer (PrimState m) Int
queMin
Bool
b2 <- Heap (PrimState m) (cost, Int) -> m Bool
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Heap (PrimState m) a -> m Bool
ACIMH.null Heap (PrimState m) (cost, Int)
heap
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not Bool
b1 Bool -> Bool -> Bool
|| Bool -> Bool
not Bool
b2) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
Int
v <-
if Bool -> Bool
not Bool
b1
then do
Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> m (Maybe Int) -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Buffer (PrimState m) Int -> m (Maybe Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Maybe a)
ACIB.popBack Buffer (PrimState m) Int
queMin
else do
(!cost
_, !Int
to) <- Maybe (cost, Int) -> (cost, Int)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (cost, Int) -> (cost, Int))
-> m (Maybe (cost, Int)) -> m (cost, Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Heap (PrimState m) (cost, Int) -> m (Maybe (cost, Int))
forall a (m :: * -> *).
(HasCallStack, Ord a, Unbox a, PrimMonad m) =>
Heap (PrimState m) a -> m (Maybe a)
ACIMH.pop Heap (PrimState m) (cost, Int)
heap
Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
to
Bit Bool
visV <- MVector (PrimState m) Bit -> Int -> Bit -> m Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector (PrimState m) Bit
vis Int
v (Bit -> m Bit) -> Bit -> m Bit
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
t) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
visV (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
cost
dualV <- MVector (PrimState m) cost -> Int -> m cost
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) cost
duals Int
v
cost
distV <- MVector (PrimState m) cost -> Int -> m cost
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) cost
dists Int
v
let start :: Int
start = Vector Int
startCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v
Vector (Int, Int, cost)
-> (Int -> (Int, Int, cost) -> m ()) -> m ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ (Csr (PrimState m) cap cost -> Int -> Vector (Int, Int, cost)
forall cap cost s.
(HasCallStack, Num cap, Unbox cap, Unbox cost) =>
Csr s cap cost -> Int -> Vector (Int, Int, cost)
ACIMCSR.adj Csr (PrimState m) cap cost
csr Int
v) ((Int -> (Int, Int, cost) -> m ()) -> m ())
-> (Int -> (Int, Int, cost) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
di (!Int
to, !Int
rev, !cost
cost) -> do
cap
cap <- MVector (PrimState m) cap -> Int -> m cap
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) cap
capCsr (Int -> m cap) -> Int -> m cap
forall a b. (a -> b) -> a -> b
$ Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
di
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (cap
cap cap -> cap -> Bool
forall a. Eq a => a -> a -> Bool
== cap
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
cost
cost' <- do
cost
dualTo <- MVector (PrimState m) cost -> Int -> m cost
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) cost
duals Int
to
cost -> m cost
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (cost -> m cost) -> cost -> m cost
forall a b. (a -> b) -> a -> b
$! cost
cost cost -> cost -> cost
forall a. Num a => a -> a -> a
- cost
dualTo cost -> cost -> cost
forall a. Num a => a -> a -> a
+ cost
dualV
cost
distTo <- MVector (PrimState m) cost -> Int -> m cost
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) cost
dists Int
to
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (cost
distTo cost -> cost -> cost
forall a. Num a => a -> a -> a
- cost
distV cost -> cost -> Bool
forall a. Ord a => a -> a -> Bool
> cost
cost') (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
let !distTo' :: cost
distTo' = cost
distV cost -> cost -> cost
forall a. Num a => a -> a -> a
+ cost
cost'
MVector (PrimState m) cost -> Int -> cost -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) cost
dists Int
to cost
distTo'
MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) Int
prevE Int
to Int
rev
if cost
distTo' cost -> cost -> Bool
forall a. Eq a => a -> a -> Bool
== cost
distV
then Buffer (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
ACIB.pushBack Buffer (PrimState m) Int
queMin Int
to
else Heap (PrimState m) (cost, Int) -> (cost, Int) -> m ()
forall a (m :: * -> *).
(HasCallStack, Ord a, Unbox a, PrimMonad m) =>
Heap (PrimState m) a -> a -> m ()
ACIMH.push Heap (PrimState m) (cost, Int)
heap (cost
distTo', Int
to)
m ()
loop
Bit Bool
visT <- MVector (PrimState m) Bit -> Int -> m Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Bit
vis Int
t
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
visT (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
cost
distT <- MVector (PrimState m) cost -> Int -> m cost
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) cost
dists Int
t
Vector Bit
vis' <- MVector (PrimState m) Bit -> m (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) Bit
vis
Vector cost
dists' <- MVector (PrimState m) cost -> m (Vector cost)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) cost
dists
Vector (Bit, cost) -> (Int -> (Bit, cost) -> m ()) -> m ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ (Vector Bit -> Vector cost -> Vector (Bit, cost)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Bit
vis' Vector cost
dists') ((Int -> (Bit, cost) -> m ()) -> m ())
-> (Int -> (Bit, cost) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
v (Bit !Bool
visV, !cost
distV) -> do
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
visV (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
MVector (PrimState m) cost -> (cost -> cost) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState m) cost
duals (cost -> cost -> cost
forall a. Num a => a -> a -> a
subtract (cost
distT cost -> cost -> cost
forall a. Num a => a -> a -> a
- cost
distV)) Int
v
Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
visT
GrowVec (PrimState m) (cap, cost)
result <- Int -> m (GrowVec (PrimState m) (cap, cost))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (GrowVec (PrimState m) a)
ACIGV.new Int
16
GrowVec (PrimState m) (cap, cost) -> (cap, cost) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
ACIGV.pushBack GrowVec (PrimState m) (cap, cost)
result (cap
0 :: cap, cost
0 :: cost)
let inner :: cap -> cost -> cost -> m ()
inner :: cap -> cost -> cost -> m ()
inner cap
flow_ cost
cost cost
prevCostPerFlow =
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (cap
flow_ cap -> cap -> Bool
forall a. Ord a => a -> a -> Bool
< cap
flowLimit) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
Bool
b <- m Bool
dualRef
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
Vector Int
prevE' <- MVector (PrimState m) Int -> m (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) Int
prevE
let minC :: cap -> Int -> m cap
minC :: cap -> Int -> m cap
minC !cap
acc Int
v
| Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
s = cap -> m cap
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure cap
acc
| Bool
otherwise = do
let iPrev :: Int
iPrev = Vector Int
prevE' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v
cap
cap <- MVector (PrimState m) cap -> Int -> m cap
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) cap
capCsr (Int -> m cap) -> Int -> m cap
forall a b. (a -> b) -> a -> b
$ Vector Int
revCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
iPrev
cap -> Int -> m cap
minC (cap -> cap -> cap
forall a. Ord a => a -> a -> a
min cap
acc cap
cap) (Int -> m cap) -> Int -> m cap
forall a b. (a -> b) -> a -> b
$ Vector Int
toCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
iPrev
cap
c <- cap -> Int -> m cap
minC (cap
flowLimit cap -> cap -> cap
forall a. Num a => a -> a -> a
- cap
flow_) Int
t
let subC :: Int -> m ()
subC :: Int -> m ()
subC Int
v = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
s) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
let iPrev :: Int
iPrev = Vector Int
prevE' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v
MVector (PrimState m) cap -> (cap -> cap) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState m) cap
capCsr (cap -> cap -> cap
forall a. Num a => a -> a -> a
+ cap
c) Int
iPrev
MVector (PrimState m) cap -> (cap -> cap) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState m) cap
capCsr (cap -> cap -> cap
forall a. Num a => a -> a -> a
subtract cap
c) (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Vector Int
revCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
iPrev
Int -> m ()
subC (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Vector Int
toCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
iPrev
Int -> m ()
subC Int
t
cost
d <- cost -> cost
forall a. Num a => a -> a
negate (cost -> cost) -> m cost -> m cost
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) cost -> Int -> m cost
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) cost
duals Int
s
let !flow' :: cap
flow' = cap
flow_ cap -> cap -> cap
forall a. Num a => a -> a -> a
+ cap
c
let !cost' :: cost
cost' = cost
cost cost -> cost -> cost
forall a. Num a => a -> a -> a
+ cap -> cost
forall a b. (Integral a, Num b) => a -> b
fromIntegral cap
c cost -> cost -> cost
forall a. Num a => a -> a -> a
* cost
d
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (cost
prevCostPerFlow cost -> cost -> Bool
forall a. Eq a => a -> a -> Bool
== cost
d) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
GrowVec (PrimState m) (cap, cost) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m ()
ACIGV.popBack_ GrowVec (PrimState m) (cap, cost)
result
GrowVec (PrimState m) (cap, cost) -> (cap, cost) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
ACIGV.pushBack GrowVec (PrimState m) (cap, cost)
result (cap
flow', cost
cost')
cap -> cost -> cost -> m ()
inner cap
flow' cost
cost' cost
d
cap -> cost -> cost -> m ()
inner cap
0 cost
0 (-cost
1)
GrowVec (PrimState m) (cap, cost) -> m (Vector (cap, cost))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
ACIGV.unsafeFreeze GrowVec (PrimState m) (cap, cost)
result
{-# INLINE getEdge #-}
getEdge ::
(HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
McfGraph (PrimState m) cap cost ->
Int ->
m (Int, Int, cap, cap, cost)
getEdge :: forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Num cap, Ord cap, Unbox cap, Num cost,
Ord cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> Int -> m (Int, Int, cap, cap, cost)
getEdge McfGraph {Int
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
nG :: forall s cap cost. McfGraph s cap cost -> Int
edgesG :: forall s cap cost.
McfGraph s cap cost -> GrowVec s (Int, Int, cap, cap, cost)
nG :: Int
edgesG :: GrowVec (PrimState m) (Int, Int, cap, cap, cost)
..} Int
i = do
Int
m <- GrowVec (PrimState m) (Int, Int, cap, cap, cost) -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m Int
ACIGV.length GrowVec (PrimState m) (Int, Int, cap, cap, cost)
edgesG
let !()
_ = HasCallStack => String -> Int -> Int -> ()
String -> Int -> Int -> ()
ACIA.checkEdge String
"AtCoder.MinCostFlow.getEdge" Int
i Int
m
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
-> Int -> m (Int, Int, cap, cap, cost)
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> Int -> m a
ACIGV.read GrowVec (PrimState m) (Int, Int, cap, cap, cost)
edgesG Int
i
{-# INLINE edges #-}
edges ::
(HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
McfGraph (PrimState m) cap cost ->
m (VU.Vector (Int, Int, cap, cap, cost))
edges :: forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Num cap, Ord cap, Unbox cap, Num cost,
Ord cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> m (Vector (Int, Int, cap, cap, cost))
edges McfGraph {Int
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
nG :: forall s cap cost. McfGraph s cap cost -> Int
edgesG :: forall s cap cost.
McfGraph s cap cost -> GrowVec s (Int, Int, cap, cap, cost)
nG :: Int
edgesG :: GrowVec (PrimState m) (Int, Int, cap, cap, cost)
..} = do
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
-> m (Vector (Int, Int, cap, cap, cost))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
ACIGV.freeze GrowVec (PrimState m) (Int, Int, cap, cap, cost)
edgesG
{-# INLINE unsafeEdges #-}
unsafeEdges ::
(HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>
McfGraph (PrimState m) cap cost ->
m (VU.Vector (Int, Int, cap, cap, cost))
unsafeEdges :: forall (m :: * -> *) cap cost.
(HasCallStack, PrimMonad m, Num cap, Ord cap, Unbox cap, Num cost,
Ord cost, Unbox cost) =>
McfGraph (PrimState m) cap cost
-> m (Vector (Int, Int, cap, cap, cost))
unsafeEdges McfGraph {Int
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
nG :: forall s cap cost. McfGraph s cap cost -> Int
edgesG :: forall s cap cost.
McfGraph s cap cost -> GrowVec s (Int, Int, cap, cap, cost)
nG :: Int
edgesG :: GrowVec (PrimState m) (Int, Int, cap, cap, cost)
..} = do
GrowVec (PrimState m) (Int, Int, cap, cap, cost)
-> m (Vector (Int, Int, cap, cap, cost))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
ACIGV.unsafeFreeze GrowVec (PrimState m) (Int, Int, cap, cap, cost)
edgesG