{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE RecordWildCards #-}

-- | Re-export of the @Csr@ module and additional graph search functions.
--
-- @since 1.1.0.0
module AtCoder.Extra.Graph
  ( -- * Re-export of CSR

    -- | The `Csr.Csr` data type and all the functions such as `build` or `adj` are re-exported.
    -- See the @Csr@ module for details.
    module Csr,

    -- * CSR helpers
    swapDupe,
    swapDupe',
    scc,
    rev,

    -- * Generic graph functions
    topSort,
    connectedComponents,
    bipartiteVertexColors,
    blockCut,
    blockCutComponents,

    -- * Shortest path search

    -- | Most of the functions are opinionated as the followings:
    --
    -- - Indices are abstracted with `Ix0` (n-dimensional `Int`).
    -- - Functions that return a predecessor array are named as @tracking*@.

    -- ** BFS (breadth-first search)

    -- *** Constraints

    -- | - Edge weight \(w > 0\)
    bfs,
    trackingBfs,

    -- ** 01-BFS

    -- *** Constraints

    -- | - Edge weight \(w\) is either \(0\) or \(1\) of type `Int`.
    bfs01,
    trackingBfs01,

    -- ** Dijkstra's algorithm

    -- *** Constraints

    -- | - Edge weight \(w > 0\)
    dijkstra,
    trackingDijkstra,

    -- ** Bellman–ford algorithm

    -- | - Vertex type is restricted to one-dimensional `Int`.
    bellmanFord,
    trackingBellmanFord,

    -- ** Floyd–Warshall algorithm (all-pair shortest path)
    floydWarshall,
    trackingFloydWarshall,

    -- *** Incremental Floyd–Warshall algorithm
    newFloydWarshall,
    newTrackingFloydWarshall,
    updateEdgeFloydWarshall,
    updateEdgeTrackingFloydWarshall,

    -- ** Path reconstruction

    -- *** Single start point (root)

    -- | Functions for retrieving a path from a predecessor array where @-1@ represents none.
    constructPathFromRoot,
    constructPathToRoot,

    -- *** All-pair

    -- | Functions for retrieving a path from a predecessor matrix \(m\), which is accessed as
    -- @m VG.! (n * from + to)@, where @-1@ represents none.
    constructPathFromRootMat,
    constructPathToRootMat,
    constructPathFromRootMatM,
    constructPathToRootMatM,
  )
where

import AtCoder.Dsu qualified as Dsu
import AtCoder.Extra.IntSet qualified as IS
import AtCoder.Extra.Ix0 (Bounds0, Ix0 (..))
import AtCoder.Internal.Buffer qualified as B
import AtCoder.Internal.Csr as Csr
import AtCoder.Internal.MinHeap qualified as MH
import AtCoder.Internal.Queue qualified as Q
import AtCoder.Internal.Scc qualified as ACISCC
import Control.Monad (replicateM_, when)
import Control.Monad.Fix (fix)
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST, runST)
import Data.Bit (Bit (..))
import Data.Foldable (for_)
import Data.Maybe (fromJust)
import Data.Vector qualified as V
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.Mutable qualified as VUM
import GHC.Stack (HasCallStack)

-- | \(O(n)\) Converts directed edges into non-directed edges; each edge \((u, v, w)\) is duplicated
-- to be \((u, v, w)\) and \((v, u, w)\). This is a convenient function for making an input to
-- `build`.
--
-- ==== __Example__
-- `swapDupe` duplicates each edge reversing the direction:
--
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> Gr.swapDupe $ VU.fromList [(0, 1, ()), (1, 2, ())]
-- [(0,1,()),(1,0,()),(1,2,()),(2,1,())]
--
-- Create a non-directed graph:
--
-- >>> let gr = Gr.build 3 . Gr.swapDupe $ VU.fromList [(0, 1, ()), (1, 2, ())]
-- >>> gr `Gr.adj` 0
-- [1]
--
-- >>> gr `Gr.adj` 1
-- [0,2]
--
-- >>> gr `Gr.adj` 2
-- [1]
--
-- @since 1.1.0.0
{-# INLINEABLE swapDupe #-}
swapDupe :: (VU.Unbox w) => VU.Vector (Int, Int, w) -> VU.Vector (Int, Int, w)
swapDupe :: forall w. Unbox w => Vector (Int, Int, w) -> Vector (Int, Int, w)
swapDupe Vector (Int, Int, w)
uvws = (forall s. ST s (MVector s (Int, Int, w))) -> Vector (Int, Int, w)
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s (Int, Int, w)))
 -> Vector (Int, Int, w))
-> (forall s. ST s (MVector s (Int, Int, w)))
-> Vector (Int, Int, w)
forall a b. (a -> b) -> a -> b
$ do
  MVector s (Int, Int, w)
vec <- Int -> ST s (MVector (PrimState (ST s)) (Int, Int, w))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Vector (Int, Int, w) -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector (Int, Int, w)
uvws)
  Vector (Int, Int, w)
-> (Int -> (Int, Int, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector (Int, Int, w)
uvws ((Int -> (Int, Int, w) -> ST s ()) -> ST s ())
-> (Int -> (Int, Int, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i (!Int
u, !Int
v, !w
w) -> do
    MVector (PrimState (ST s)) (Int, Int, w)
-> Int -> (Int, Int, w) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s (Int, Int, w)
MVector (PrimState (ST s)) (Int, Int, w)
vec (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) (Int
u, Int
v, w
w)
    MVector (PrimState (ST s)) (Int, Int, w)
-> Int -> (Int, Int, w) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s (Int, Int, w)
MVector (PrimState (ST s)) (Int, Int, w)
vec (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
v, Int
u, w
w)
  MVector s (Int, Int, w) -> ST s (MVector s (Int, Int, w))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s (Int, Int, w)
vec

-- | \(O(n)\) Converts directed edges into non-directed edges; each edge \((u, v)\) is duplicated
-- to be \((u, v)\) and \((v, u)\). This is a convenient function for making an input to `build'`.
--
-- ==== __Example__
-- `swapDupe'` duplicates each edge reversing the direction:
--
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2)]
-- [(0,1),(1,0),(1,2),(2,1)]
--
-- Create a non-directed graph:
--
-- >>> let gr = Gr.build' 3 . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2)]
-- >>> gr `Gr.adj` 0
-- [1]
--
-- >>> gr `Gr.adj` 1
-- [0,2]
--
-- >>> gr `Gr.adj` 2
-- [1]
--
-- @since 1.1.0.0
{-# INLINEABLE swapDupe' #-}
-- NOTE: concatMap does not fuse anyways, as the vector's code says
swapDupe' :: VU.Vector (Int, Int) -> VU.Vector (Int, Int)
swapDupe' :: Vector (Int, Int) -> Vector (Int, Int)
swapDupe' Vector (Int, Int)
uvs = (forall s. ST s (MVector s (Int, Int))) -> Vector (Int, Int)
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s (Int, Int))) -> Vector (Int, Int))
-> (forall s. ST s (MVector s (Int, Int))) -> Vector (Int, Int)
forall a b. (a -> b) -> a -> b
$ do
  MVector s (Int, Int)
vec <- Int -> ST s (MVector (PrimState (ST s)) (Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Vector (Int, Int) -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector (Int, Int)
uvs)
  Vector (Int, Int) -> (Int -> (Int, Int) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector (Int, Int)
uvs ((Int -> (Int, Int) -> ST s ()) -> ST s ())
-> (Int -> (Int, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i (!Int
u, !Int
v) -> do
    MVector (PrimState (ST s)) (Int, Int)
-> Int -> (Int, Int) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s (Int, Int)
MVector (PrimState (ST s)) (Int, Int)
vec (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) (Int
u, Int
v)
    MVector (PrimState (ST s)) (Int, Int)
-> Int -> (Int, Int) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s (Int, Int)
MVector (PrimState (ST s)) (Int, Int)
vec (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
v, Int
u)
  MVector s (Int, Int) -> ST s (MVector s (Int, Int))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s (Int, Int)
vec

-- | \(O(n + m)\) Returns the strongly connected components of a `Csr`.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> -- 0 == 1 -> 2    3
-- >>> let gr = Gr.build' 4 $ VU.fromList [(0, 1), (1, 0), (1, 2)]
-- >>> Gr.scc gr
-- [[3],[0,1],[2]]
--
-- @since 1.1.0.0
{-# INLINE scc #-}
scc :: Csr w -> V.Vector (VU.Vector Int)
scc :: forall w. Csr w -> Vector (Vector Int)
scc = Csr w -> Vector (Vector Int)
forall w. Csr w -> Vector (Vector Int)
ACISCC.sccCsr

-- | \(O(n + m)\) Returns a reverse graph, where original edges \((u, v, w)\) are transposed to be
-- \((v, u, w)\). Reverse graphs are useful for, for example, getting distance to a specific vertex
-- from every other vertex with `dijkstra`.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> -- 0 == 1 -> 2 -> 3
-- >>> let gr = Gr.build' 4 $ VU.fromList [(0, 1), (1, 0), (1, 2), (2, 3)]
-- >>> map (Gr.adj gr) [0 .. 3]
-- [[1],[0,2],[3],[]]
--
-- >>> -- 0 == 1 <- 2 <- 3
-- >>> let revGr = Gr.rev gr
-- >>> map (Gr.adj revGr) [0 .. 3]
-- [[1],[0],[1],[2]]
--
-- @since 1.2.3.0
{-# INLINEABLE rev #-}
rev :: (VU.Unbox w) => Csr w -> Csr w
rev :: forall w. Unbox w => Csr w -> Csr w
rev Csr {Int
Vector w
Vector Int
nCsr :: Int
mCsr :: Int
startCsr :: Vector Int
adjCsr :: Vector Int
wCsr :: Vector w
nCsr :: forall w. Csr w -> Int
mCsr :: forall w. Csr w -> Int
startCsr :: forall w. Csr w -> Vector Int
adjCsr :: forall w. Csr w -> Vector Int
wCsr :: forall w. Csr w -> Vector w
..} = Int -> Vector (Int, Int, w) -> Csr w
forall w.
(HasCallStack, Unbox w) =>
Int -> Vector (Int, Int, w) -> Csr w
Csr.build Int
nCsr Vector (Int, Int, w)
revEdges
  where
    vws :: Vector (Int, w)
vws = Vector Int -> Vector w -> Vector (Int, w)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Int
adjCsr Vector w
wCsr
    revEdges :: Vector (Int, Int, w)
revEdges = ((Int -> Vector (Int, Int, w))
 -> Vector Int -> Vector (Int, Int, w))
-> Vector Int
-> (Int -> Vector (Int, Int, w))
-> Vector (Int, Int, w)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int -> Vector (Int, Int, w)) -> Vector Int -> Vector (Int, Int, w)
forall a b.
(Unbox a, Unbox b) =>
(a -> Vector b) -> Vector a -> Vector b
VU.concatMap (Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
nCsr Int -> Int
forall a. a -> a
id) ((Int -> Vector (Int, Int, w)) -> Vector (Int, Int, w))
-> (Int -> Vector (Int, Int, w)) -> Vector (Int, Int, w)
forall a b. (a -> b) -> a -> b
$ \Int
v1 ->
      let !o1 :: Int
o1 = Vector Int
startCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v1
          !o2 :: Int
o2 = Vector Int
startCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
v1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          !vw2s :: Vector (Int, w)
vw2s = Int -> Int -> Vector (Int, w) -> Vector (Int, w)
forall a. Unbox a => Int -> Int -> Vector a -> Vector a
VU.slice Int
o1 (Int
o2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
o1) Vector (Int, w)
vws
       in ((Int, w) -> (Int, Int, w))
-> Vector (Int, w) -> Vector (Int, Int, w)
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (\(!Int
v2, !w
w2) -> (Int
v2, Int
v1, w
w2)) Vector (Int, w)
vw2s

-- -------------------------------------------------------------------------------------------------
-- Generic graph search functions
-- -------------------------------------------------------------------------------------------------

-- | \(O(n \log n + m)\) Returns the lexicographically smallest topological ordering of the given
-- graph.
--
-- ==== Constraints
-- - The graph must be a DAG; no cycle can exist.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let n = 5
-- >>> let gr = Gr.build' n $ VU.fromList [(1, 2), (4, 0), (0, 3)]
-- >>> Gr.topSort n (gr `Gr.adj`)
-- [1,2,4,0,3]
--
-- @since 1.1.0.0
{-# INLINEABLE topSort #-}
topSort :: Int -> (Int -> VU.Vector Int) -> VU.Vector Int
topSort :: Int -> (Int -> Vector Int) -> Vector Int
topSort Int
n Int -> Vector Int
gr = (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int)) -> Vector Int)
-> (forall s. ST s (Vector Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
  MVector s Int
inDeg <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
n (Int
0 :: Int)
  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
u -> do
    Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Int -> Vector Int
gr Int
u) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v -> do
      MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
inDeg (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
v

  -- start from the vertices with zero in-degrees:
  IntSet s
que <- Int -> ST s (IntSet (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (IntSet (PrimState m))
IS.new Int
n
  Vector Int
inDeg' <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
inDeg
  Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
inDeg' ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v Int
d -> do
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      IntSet (PrimState (ST s)) -> Int -> ST s ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
IntSet (PrimState m) -> Int -> m ()
IS.insert IntSet s
IntSet (PrimState (ST s))
que Int
v

  Buffer s Int
buf <- Int -> ST s (Buffer (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Buffer (PrimState m) a)
B.new Int
n
  (ST s () -> ST s ()) -> ST s ()
forall a. (a -> a) -> a
fix ((ST s () -> ST s ()) -> ST s ())
-> (ST s () -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ST s ()
loop -> do
    IntSet (PrimState (ST s)) -> ST s (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> m (Maybe Int)
IS.deleteMin IntSet s
IntSet (PrimState (ST s))
que ST s (Maybe Int) -> (Maybe Int -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Maybe Int
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Just Int
u -> do
        Buffer (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer s Int
Buffer (PrimState (ST s)) Int
buf Int
u
        Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Int -> Vector Int
gr Int
u) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v -> do
          Int
nv <- Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1 (Int -> Int) -> ST s Int -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
inDeg Int
v
          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
inDeg Int
v Int
nv
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
nv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            IntSet (PrimState (ST s)) -> Int -> ST s ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
IntSet (PrimState m) -> Int -> m ()
IS.insert IntSet s
IntSet (PrimState (ST s))
que Int
v
        ST s ()
loop

  Buffer (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Vector a)
B.unsafeFreeze Buffer s Int
Buffer (PrimState (ST s)) Int
buf

-- | \(O(n)\) Returns connected components for a non-directed graph.
--
-- ==== Constraints
-- - The graph must be non-directed.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1), (1, 2)]
-- >>> let gr = Gr.build' 4 $ Gr.swapDupe' es
-- >>> Gr.connectedComponents 4 (Gr.adj gr)
-- [[0,1,2],[3]]
--
-- >>> Gr.connectedComponents 0 (const VU.empty)
-- []
--
-- @since 1.2.4.0
{-# INLINEABLE connectedComponents #-}
connectedComponents :: Int -> (Int -> VU.Vector Int) -> V.Vector (VU.Vector Int)
connectedComponents :: Int -> (Int -> Vector Int) -> Vector (Vector Int)
connectedComponents Int
n Int -> Vector Int
gr = (forall s. ST s (Vector (Vector Int))) -> Vector (Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector (Vector Int))) -> Vector (Vector Int))
-> (forall s. ST s (Vector (Vector Int))) -> Vector (Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  Buffer (PrimState (ST s)) Int
buf <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Buffer (PrimState m) a)
B.new @_ @Int Int
n
  Buffer s Int
len <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Buffer (PrimState m) a)
B.new @_ @Int Int
n
  MVector (PrimState (ST s)) Bit
vis <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Bit Int
n (Bool -> Bit
Bit Bool
False)

  let dfs :: b -> Int -> ST s b
dfs !b
acc Int
u = do
        Bit Bool
b <- MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector (PrimState (ST s)) Bit
vis Int
u (Bit -> ST s Bit) -> Bit -> ST s Bit
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
        if Bool
b
          then b -> ST s b
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
acc
          else do
            Buffer (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer (PrimState (ST s)) Int
buf Int
u
            (b -> Int -> ST s b) -> b -> Vector Int -> ST s b
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM' b -> Int -> ST s b
dfs (b
acc b -> b -> b
forall a. Num a => a -> a -> a
+ b
1) (Int -> Vector Int
gr Int
u)

  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
u -> do
    Int
l :: Int <- Int -> Int -> ST s Int
forall {b}. Num b => b -> Int -> ST s b
dfs Int
0 Int
u
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      Buffer (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer s Int
Buffer (PrimState (ST s)) Int
len Int
l

  Vector Int
vs0 <- Buffer (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Vector a)
B.unsafeFreeze Buffer (PrimState (ST s)) Int
buf
  Vector Int
lens0 <- Buffer (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Vector a)
B.unsafeFreeze Buffer s Int
Buffer (PrimState (ST s)) Int
len

  Vector (Vector Int) -> ST s (Vector (Vector Int))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    (Vector (Vector Int) -> ST s (Vector (Vector Int)))
-> ((Vector Int, Vector Int) -> Vector (Vector Int))
-> (Vector Int, Vector Int)
-> ST s (Vector (Vector Int))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> ((Vector Int, Vector Int)
    -> (Vector Int, (Vector Int, Vector Int)))
-> (Vector Int, Vector Int)
-> Vector (Vector Int)
forall b a. Int -> (b -> (a, b)) -> b -> Vector a
V.unfoldrExactN
      (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
lens0)
      ( \(!Vector Int
vs, !Vector Int
ls) ->
          let (!Int
l, !Vector Int
lsR) = Maybe (Int, Vector Int) -> (Int, Vector Int)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Int, Vector Int) -> (Int, Vector Int))
-> Maybe (Int, Vector Int) -> (Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Maybe (Int, Vector Int)
forall a. Unbox a => Vector a -> Maybe (a, Vector a)
VU.uncons Vector Int
ls
              (!Vector Int
vsL, !Vector Int
vsR) = Int -> Vector Int -> (Vector Int, Vector Int)
forall a. Unbox a => Int -> Vector a -> (Vector a, Vector a)
VU.splitAt Int
l Vector Int
vs
           in (Vector Int
vsL, (Vector Int
vsR, Vector Int
lsR))
      )
    ((Vector Int, Vector Int) -> ST s (Vector (Vector Int)))
-> (Vector Int, Vector Int) -> ST s (Vector (Vector Int))
forall a b. (a -> b) -> a -> b
$ (Vector Int
vs0, Vector Int
lens0)

-- | \(O((n + m) \alpha)\) Returns a bipartite vertex coloring for a bipartite graph.
-- Returns `Nothing` for a non-bipartite graph.
--
-- ==== Constraints
-- - The graph must not be directed.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1), (1, 2)]
-- >>> let gr = Gr.build' 4 es
-- >>> Gr.bipartiteVertexColors 4 (Gr.adj gr)
-- Just [0,1,0,0]
--
-- @since 1.2.4.0
{-# INLINEABLE bipartiteVertexColors #-}
bipartiteVertexColors :: Int -> (Int -> VU.Vector Int) -> Maybe (VU.Vector Bit)
bipartiteVertexColors :: Int -> (Int -> Vector Int) -> Maybe (Vector Bit)
bipartiteVertexColors Int
n Int -> Vector Int
gr = (forall s. ST s (Maybe (Vector Bit))) -> Maybe (Vector Bit)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe (Vector Bit))) -> Maybe (Vector Bit))
-> (forall s. ST s (Maybe (Vector Bit))) -> Maybe (Vector Bit)
forall a b. (a -> b) -> a -> b
$ do
  (!Bool
isBipartite, !Vector Bit
color, !Dsu s
_) <- Int -> (Int -> Vector Int) -> ST s (Bool, Vector Bit, Dsu s)
forall s.
Int -> (Int -> Vector Int) -> ST s (Bool, Vector Bit, Dsu s)
bipartiteVertexColorsImpl Int
n Int -> Vector Int
gr
  if Bool
isBipartite
    then Maybe (Vector Bit) -> ST s (Maybe (Vector Bit))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Vector Bit) -> ST s (Maybe (Vector Bit)))
-> Maybe (Vector Bit) -> ST s (Maybe (Vector Bit))
forall a b. (a -> b) -> a -> b
$ Vector Bit -> Maybe (Vector Bit)
forall a. a -> Maybe a
Just Vector Bit
color
    else Maybe (Vector Bit) -> ST s (Maybe (Vector Bit))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Vector Bit)
forall a. Maybe a
Nothing

{-# INLINEABLE bipartiteVertexColorsImpl #-}
bipartiteVertexColorsImpl :: Int -> (Int -> VU.Vector Int) -> ST s (Bool, VU.Vector Bit, Dsu.Dsu s)
bipartiteVertexColorsImpl :: forall s.
Int -> (Int -> Vector Int) -> ST s (Bool, Vector Bit, Dsu s)
bipartiteVertexColorsImpl Int
n Int -> Vector Int
gr
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = do
      Dsu s
dsu <- Int -> ST s (Dsu (PrimState (ST s)))
forall (m :: * -> *). PrimMonad m => Int -> m (Dsu (PrimState m))
Dsu.new Int
0
      (Bool, Vector Bit, Dsu s) -> ST s (Bool, Vector Bit, Dsu s)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
True, Vector Bit
forall a. Unbox a => Vector a
VU.empty, Dsu s
dsu)
  | Bool
otherwise = do
      -- 0 <= v < n: red, n <= v: green
      Dsu (PrimState (ST s))
dsu <- Int -> ST s (Dsu (PrimState (ST s)))
forall (m :: * -> *). PrimMonad m => Int -> m (Dsu (PrimState m))
Dsu.new (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n)
      [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
u -> do
        Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Int -> Vector Int
gr Int
u) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v -> do
          -- try both (red, green) and (green, red) colorings:
          Dsu (PrimState (ST s)) -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
Dsu (PrimState m) -> Int -> Int -> m ()
Dsu.merge_ Dsu (PrimState (ST s))
dsu (Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) Int
v
          Dsu (PrimState (ST s)) -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
Dsu (PrimState m) -> Int -> Int -> m ()
Dsu.merge_ Dsu (PrimState (ST s))
dsu Int
u (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)

      MVector (PrimState (ST s)) Bit
color <- Int -> Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) (Bit -> ST s (MVector (PrimState (ST s)) Bit))
-> Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False

      -- for each leader vertices, paint their colors:
      [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v -> do
        Int
l <- Dsu (PrimState (ST s)) -> Int -> ST s Int
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
Dsu (PrimState m) -> Int -> m Int
Dsu.leader Dsu (PrimState (ST s))
dsu Int
v
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
v) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Bit
color (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) (Bit -> ST s ()) -> Bit -> ST s ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True

      -- paint other vertices:
      [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v -> do
        MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Bit
color Int
v (Bit -> ST s ()) -> ST s Bit -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Bit
color (Int -> ST s Bit) -> ST s Int -> ST s Bit
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Dsu (PrimState (ST s)) -> Int -> ST s Int
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
Dsu (PrimState m) -> Int -> m Int
Dsu.leader Dsu (PrimState (ST s))
dsu Int
v
        MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Bit
color (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) (Bit -> ST s ()) -> ST s Bit -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Bit
color (Int -> ST s Bit) -> ST s Int -> ST s Bit
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Dsu (PrimState (ST s)) -> Int -> ST s Int
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
Dsu (PrimState m) -> Int -> m Int
Dsu.leader Dsu (PrimState (ST s))
dsu (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)

      Vector Bit
color' <- MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (MVector (PrimState (ST s)) Bit -> ST s (Vector Bit))
-> MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a b. (a -> b) -> a -> b
$ Int
-> MVector (PrimState (ST s)) Bit -> MVector (PrimState (ST s)) Bit
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VGM.take Int
n MVector (PrimState (ST s)) Bit
color
      let isCompatible :: Int -> ST s Bool
isCompatible Int
v
            | Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Bool -> ST s Bool
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
            | Bool
otherwise = do
                Bit
c1 <- MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Bit
color (Int -> ST s Bit) -> ST s Int -> ST s Bit
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Dsu (PrimState (ST s)) -> Int -> ST s Int
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
Dsu (PrimState m) -> Int -> m Int
Dsu.leader Dsu (PrimState (ST s))
dsu Int
v
                Bit
c2 <- MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Bit
color (Int -> ST s Bit) -> ST s Int -> ST s Bit
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Dsu (PrimState (ST s)) -> Int -> ST s Int
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
Dsu (PrimState m) -> Int -> m Int
Dsu.leader Dsu (PrimState (ST s))
dsu (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
                if Bit
c1 Bit -> Bit -> Bool
forall a. Eq a => a -> a -> Bool
== Bit
c2
                  then Bool -> ST s Bool
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
                  else Int -> ST s Bool
isCompatible (Int -> ST s Bool) -> Int -> ST s Bool
forall a b. (a -> b) -> a -> b
$ Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

      Bool
b <- Int -> ST s Bool
isCompatible Int
0
      (Bool, Vector Bit, Dsu s) -> ST s (Bool, Vector Bit, Dsu s)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
b, Vector Bit
color', Dsu s
Dsu (PrimState (ST s))
dsu)

-- | \(O(n + m)\) Returns a [block cut tree](https://en.wikipedia.org/wiki/Biconnected_component)
-- where super vertices \((v \ge n)\) represent each biconnected component.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> -- 0---3---2
-- >>> -- +-1-+
-- >>> let n = 4
-- >>> let gr = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 3), (0, 1), (1, 3), (3, 2)]
-- >>> let bct = blockCut n (gr `Gr.adj`)
-- >>> bct
-- Csr {nCsr = 6, mCsr = 5, startCsr = [0,0,0,0,0,2,5], adjCsr = [3,2,0,3,1], wCsr = [(),(),(),(),()]}
--
-- >>> V.generate (Gr.nCsr bct - n) ((bct `Gr.adj`) . (+ n))
-- [[3,2],[0,3,1]]
--
-- @since 1.1.1.0
{-# INLINEABLE blockCut #-}
blockCut :: Int -> (Int -> VU.Vector Int) -> Csr ()
blockCut :: Int -> (Int -> Vector Int) -> Csr ()
blockCut Int
n Int -> Vector Int
gr = (forall s. ST s (Csr ())) -> Csr ()
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Csr ())) -> Csr ())
-> (forall s. ST s (Csr ())) -> Csr ()
forall a b. (a -> b) -> a -> b
$ do
  MVector (PrimState (ST s)) Int
low <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
n (Int
0 :: Int)
  MVector (PrimState (ST s)) Int
ord <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
n (Int
0 :: Int)
  Buffer (PrimState (ST s)) Int
st <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Buffer (PrimState m) a)
B.new @_ @Int Int
n
  MVector (PrimState (ST s)) Bit
used <- Int -> Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
n (Bit -> ST s (MVector (PrimState (ST s)) Bit))
-> Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False
  Buffer (PrimState (ST s)) (Int, Int)
edges <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Buffer (PrimState m) a)
B.new @_ @(Int, Int {- TODO: correct capacity? -}) (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n)
  -- represents the bidirected component's index. also works as super vertex indices.
  MVector (PrimState (ST s)) Int
next <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 Int
n

  let dfs :: Int -> Int -> Int -> ST s Int
dfs Int
k0 Int
v Int
p = do
        Buffer (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer (PrimState (ST s)) Int
st Int
v
        MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Bit
used Int
v (Bit -> ST s ()) -> Bit -> ST s ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
        MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
low Int
v Int
k0
        MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
ord Int
v Int
k0

        (Int, Int) -> Int
forall a b. (a, b) -> b
snd
          ((Int, Int) -> Int) -> ST s (Int, Int) -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Int, Int) -> Int -> ST s (Int, Int))
-> (Int, Int) -> Vector Int -> ST s (Int, Int)
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM'
            ( \(!Int
child, !Int
k) Int
to -> do
                if Int
to Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
p
                  then (Int, Int) -> ST s (Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
child, Int
k)
                  else do
                    Bit Bool
b <- MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Bit
used Int
to
                    if Bool -> Bool
not Bool
b
                      then do
                        let !child' :: Int
child' = Int
child Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                        Int
s <- Buffer (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m Int
B.length Buffer (PrimState (ST s)) Int
st
                        Int
k' <- Int -> Int -> Int -> ST s Int
dfs Int
k Int
to Int
v
                        Int
lowTo <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Int
low Int
to
                        MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState (ST s)) Int
low (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
lowTo) Int
v
                        Int
ordV <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Int
ord Int
v
                        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1 Bool -> Bool -> Bool
&& Int
child' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) Bool -> Bool -> Bool
|| (Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= -Int
1 Bool -> Bool -> Bool
&& Int
lowTo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
ordV)) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                          Int
nxt <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState (ST s)) Int
next Int
0
                          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector (PrimState (ST s)) Int
next Int
0 (Int
nxt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                          Buffer (PrimState (ST s)) (Int, Int) -> (Int, Int) -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer (PrimState (ST s)) (Int, Int)
edges (Int
nxt, Int
v)
                          Int
len <- Buffer (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m Int
B.length Buffer (PrimState (ST s)) Int
st
                          Int -> ST s () -> ST s ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                            Int
back <- Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> ST s (Maybe Int) -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Buffer (PrimState (ST s)) Int -> ST s (Maybe Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Maybe a)
B.popBack Buffer (PrimState (ST s)) Int
st
                            Buffer (PrimState (ST s)) (Int, Int) -> (Int, Int) -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer (PrimState (ST s)) (Int, Int)
edges (Int
nxt, Int
back)
                        (Int, Int) -> ST s (Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
child', Int
k')
                      else do
                        Int
ordTo <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Int
ord Int
to
                        MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState (ST s)) Int
low (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
ordTo) Int
v
                        (Int, Int) -> ST s (Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
child, Int
k)
            )
            (Int
0 :: Int, Int
k0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            (Int -> Vector Int
gr Int
v)

  Int
_ <-
    (Int -> Int -> Bit -> ST s Int)
-> Int -> MVector (PrimState (ST s)) Bit -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a b.
(PrimMonad m, MVector v a) =>
(b -> Int -> a -> m b) -> b -> v (PrimState m) a -> m b
VGM.ifoldM'
      ( \Int
k Int
v (Bit Bool
b) -> do
          if Bool
b
            then do
              Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
k
            else do
              Int
k' <- Int -> Int -> Int -> ST s Int
dfs Int
k Int
v (-Int
1)
              Vector Int
st' <- Buffer (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Vector a)
B.unsafeFreeze Buffer (PrimState (ST s)) Int
st
              Int
nxt <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState (ST s)) Int
next Int
0
              MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector (PrimState (ST s)) Int
next Int
0 (Int
nxt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
              Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector Int
st' ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x -> do
                Buffer (PrimState (ST s)) (Int, Int) -> (Int, Int) -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer (PrimState (ST s)) (Int, Int)
edges (Int
nxt, Int
x)
              Buffer (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m ()
B.clear Buffer (PrimState (ST s)) Int
st
              Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
k'
      )
      (Int
0 :: Int)
      MVector (PrimState (ST s)) Bit
used

  Int
n' <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState (ST s)) Int
next Int
0
  HasCallStack => Int -> Vector (Int, Int) -> Csr ()
Int -> Vector (Int, Int) -> Csr ()
Csr.build' Int
n' (Vector (Int, Int) -> Csr ())
-> ST s (Vector (Int, Int)) -> ST s (Csr ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Buffer (PrimState (ST s)) (Int, Int) -> ST s (Vector (Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Vector a)
B.unsafeFreeze Buffer (PrimState (ST s)) (Int, Int)
edges

-- | \(O(n + m)\) Returns [blocks (biconnected comopnents)](https://en.wikipedia.org/wiki/Biconnected_component)
-- of the graph.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> -- 0---3---2
-- >>> -- +-1-+
-- >>> let n = 4
-- >>> let gr = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 3), (0, 1), (1, 3), (3, 2)]
-- >>> Gr.blockCutComponents n (gr `Gr.adj`)
-- [[3,2],[0,3,1]]
--
-- @since 1.1.1.0
{-# INLINEABLE blockCutComponents #-}
blockCutComponents :: Int -> (Int -> VU.Vector Int) -> V.Vector (VU.Vector Int)
blockCutComponents :: Int -> (Int -> Vector Int) -> Vector (Vector Int)
blockCutComponents Int
n Int -> Vector Int
gr =
  let bct :: Csr ()
bct = Int -> (Int -> Vector Int) -> Csr ()
blockCut Int
n Int -> Vector Int
gr
      d :: Int
d = Csr () -> Int
forall w. Csr w -> Int
nCsr Csr ()
bct Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n
   in Int -> (Int -> Vector Int) -> Vector (Vector Int)
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
d ((Csr ()
bct `adj`) (Int -> Vector Int) -> (Int -> Int) -> Int -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n))

-- -------------------------------------------------------------------------------------------------
-- Opinionated graph search functions
-- -------------------------------------------------------------------------------------------------

-- The implementations can be a bit simpler with `whenJustM`

-- | \(O(n + m)\) Opinionated breadth-first search that returns a distance array.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 1 :: Int), (1, 2, 10)]
-- >>> let gr = Gr.build 4 es
-- >>> Gr.bfs 4 (Gr.adjW gr) (-1) (VU.singleton (0, 0))
-- [0,1,11,-1]
--
-- @since 1.2.4.0
{-# INLINE bfs #-}
bfs ::
  forall i w.
  (HasCallStack, Ix0 i, VU.Unbox i, VU.Unbox w, Num w, Eq w) =>
  -- | Zero-based vertex boundary.
  Bounds0 i ->
  -- | Graph function that takes a vertex and returns adjacent vertices with edge weights, where
  -- \(w > 0\).
  (i -> VU.Vector (i, w)) ->
  -- | Distance assignment for unreachable vertices.
  w ->
  -- | Weighted source vertices.
  VU.Vector (i, w) ->
  -- | Distance array in one-dimensional index.
  VU.Vector w
bfs :: forall i w.
(HasCallStack, Ix0 i, Unbox i, Unbox w, Num w, Eq w) =>
i -> (i -> Vector (i, w)) -> w -> Vector (i, w) -> Vector w
bfs !i
bnd0 !i -> Vector (i, w)
gr !w
undefW !Vector (i, w)
sources =
  let (!Vector w
dist, !Vector Int
_) = Bool
-> i
-> (i -> Vector (i, w))
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
forall i w.
(HasCallStack, Ix0 i, Unbox i, Unbox w, Num w, Eq w) =>
Bool
-> i
-> (i -> Vector (i, w))
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
bfsImpl Bool
False i
bnd0 i -> Vector (i, w)
gr w
undefW Vector (i, w)
sources
   in Vector w
dist

-- | \(O(n + m)\) Opinionated breadth-first search that returns a distance array and a predecessor
-- array.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 1 :: Int), (1, 2, 10)]
-- >>> let gr = Gr.build 4 es
-- >>> let (!dist, !prev) = Gr.trackingBfs 4 (Gr.adjW gr) (-1) (VU.singleton (0, 0))
-- >>> dist
-- [0,1,11,-1]
--
-- >>> Gr.constructPathFromRoot prev 2
-- [0,1,2]
--
-- @since 1.2.4.0
{-# INLINE trackingBfs #-}
trackingBfs ::
  forall i w.
  (HasCallStack, Ix0 i, VU.Unbox i, VU.Unbox w, Num w, Eq w) =>
  -- | Zero-based vertex boundary.
  Bounds0 i ->
  -- | Graph function that takes a vertex and returns adjacent vertices with edge weights, where
  -- \(w > 0\).
  (i -> VU.Vector (i, w)) ->
  -- | Distance assignment for unreachable vertices.
  w ->
  -- | Weighted source vertices.
  VU.Vector (i, w) ->
  -- | A tuple of (Distance vector in one-dimensional index, Predecessor array (@-1@ represents none)).
  (VU.Vector w, VU.Vector Int)
trackingBfs :: forall i w.
(HasCallStack, Ix0 i, Unbox i, Unbox w, Num w, Eq w) =>
i
-> (i -> Vector (i, w))
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
trackingBfs = Bool
-> i
-> (i -> Vector (i, w))
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
forall i w.
(HasCallStack, Ix0 i, Unbox i, Unbox w, Num w, Eq w) =>
Bool
-> i
-> (i -> Vector (i, w))
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
bfsImpl Bool
True

{-# INLINEABLE bfsImpl #-}
bfsImpl ::
  forall i w.
  (HasCallStack, Ix0 i, VU.Unbox i, VU.Unbox w, Num w, Eq w) =>
  Bool ->
  Bounds0 i ->
  (i -> VU.Vector (i, w)) ->
  w ->
  VU.Vector (i, w) ->
  (VU.Vector w, VU.Vector Int)
bfsImpl :: forall i w.
(HasCallStack, Ix0 i, Unbox i, Unbox w, Num w, Eq w) =>
Bool
-> i
-> (i -> Vector (i, w))
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
bfsImpl !Bool
trackPrev !i
bnd0 !i -> Vector (i, w)
gr !w
undefW !Vector (i, w)
sources
  | Vector (i, w) -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector (i, w)
sources Bool -> Bool -> Bool
&& Bool
trackPrev = (Int -> w -> Vector w
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nVerts w
undefW, Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nVerts (-Int
1))
  | Vector (i, w) -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector (i, w)
sources = (Int -> w -> Vector w
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nVerts w
undefW, Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
0 (-Int
1))
  | Bool
otherwise = (forall s. ST s (Vector w, Vector Int)) -> (Vector w, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector w, Vector Int)) -> (Vector w, Vector Int))
-> (forall s. ST s (Vector w, Vector Int))
-> (Vector w, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
      MVector s w
dist <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @w Int
nVerts w
undefW
      MVector s Int
prev <-
        if Bool
trackPrev
          then forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
nVerts (-Int
1)
          else forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
0 (-Int
1)

      -- NOTE: We only need capacity of `n`, as first appearance of vertex is always with the
      -- minimum distance.
      Queue s i
queue <- Int -> ST s (Queue (PrimState (ST s)) i)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Queue (PrimState m) a)
Q.new Int
nVerts

      -- set source values
      Vector (i, w) -> ((i, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (i, w)
sources (((i, w) -> ST s ()) -> ST s ()) -> ((i, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!i
src, !w
w0) -> do
        -- TODO: assert w1 <= w2
        let !i :: Int
i = i -> i -> Int
forall i. Ix0 i => i -> i -> Int
index0 i
bnd0 i
src
        !w
lastD <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
dist Int
i
        -- Note that duplicate inputs are pruned here:
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
lastD w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState (ST s)) w
dist Int
i w
w0
          Queue (PrimState (ST s)) i -> i -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
Q.pushBack Queue s i
Queue (PrimState (ST s)) i
queue i
src

      -- run BFS
      (ST s () -> ST s ()) -> ST s ()
forall a. (a -> a) -> a
fix ((ST s () -> ST s ()) -> ST s ())
-> (ST s () -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ST s ()
loop -> do
        Queue (PrimState (ST s)) i -> ST s (Maybe i)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m (Maybe a)
Q.popFront Queue s i
Queue (PrimState (ST s)) i
queue ST s (Maybe i) -> (Maybe i -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Maybe i
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          Just i
v1 -> do
            let !i1 :: Int
i1 = i -> i -> Int
forall i. Ix0 i => i -> i -> Int
index0 i
bnd0 i
v1
            !w
d1 <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
dist Int
i1
            Vector (i, w) -> ((i, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (i -> Vector (i, w)
gr i
v1) (((i, w) -> ST s ()) -> ST s ()) -> ((i, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!i
v2, !w
dw) -> do
              let !i2 :: Int
i2 = i -> i -> Int
forall i. Ix0 i => i -> i -> Int
index0 i
bnd0 i
v2
              !w
lastD <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
dist Int
i2
              Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
lastD w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState (ST s)) w
dist Int
i2 (w -> ST s ()) -> w -> ST s ()
forall a b. (a -> b) -> a -> b
$! w
d1 w -> w -> w
forall a. Num a => a -> a -> a
+ w
dw
                Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                  MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
prev Int
i2 Int
i1
                Queue (PrimState (ST s)) i -> i -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
Q.pushBack Queue s i
Queue (PrimState (ST s)) i
queue i
v2
            ST s ()
loop

      (,) (Vector w -> Vector Int -> (Vector w, Vector Int))
-> ST s (Vector w) -> ST s (Vector Int -> (Vector w, Vector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) w -> ST s (Vector w)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s w
MVector (PrimState (ST s)) w
dist ST s (Vector Int -> (Vector w, Vector Int))
-> ST s (Vector Int) -> ST s (Vector w, Vector Int)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
prev
  where
    !nVerts :: Int
nVerts = i -> Int
forall i. Ix0 i => i -> Int
rangeSize0 i
bnd0

-- | \(O(n + m)\) Opinionated 01-BFS that returns a distance array.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 10 :: Int), (0, 2, 0), (2, 1, 1)]
-- >>> let gr = Gr.build 4 es
-- >>> let capacity = 10
-- >>> Gr.bfs01 4 (Gr.adjW gr) capacity (VU.singleton (0, 0))
-- [0,1,0,-1]
--
-- @since 1.2.4.0
{-# INLINE bfs01 #-}
bfs01 ::
  forall i.
  (HasCallStack, Ix0 i, VU.Unbox i) =>
  -- | Zero-based index boundary.
  Bounds0 i ->
  -- | Graph function that takes the vertexand returns adjacent vertices with edge weights, where
  -- \(w > 0\).
  (i -> VU.Vector (i, Int)) ->
  -- | Capacity of deque, often the number of edges \(m\).
  Int ->
  -- | Weighted source vertices.
  VU.Vector (i, Int) ->
  -- | Distance array in one-dimensional index. Unreachable vertices are assigned distance of @-1@.
  VU.Vector Int
bfs01 :: forall i.
(HasCallStack, Ix0 i, Unbox i) =>
i -> (i -> Vector (i, Int)) -> Int -> Vector (i, Int) -> Vector Int
bfs01 !i
bnd0 !i -> Vector (i, Int)
gr !Int
capacity !Vector (i, Int)
sources =
  let (!Vector Int
dist, !Vector Int
_) = Bool
-> i
-> (i -> Vector (i, Int))
-> Int
-> Vector (i, Int)
-> (Vector Int, Vector Int)
forall i.
(HasCallStack, Ix0 i, Unbox i) =>
Bool
-> i
-> (i -> Vector (i, Int))
-> Int
-> Vector (i, Int)
-> (Vector Int, Vector Int)
bfs01Impl Bool
False i
bnd0 i -> Vector (i, Int)
gr Int
capacity Vector (i, Int)
sources
   in Vector Int
dist

-- | \(O(n + m)\) Opinionated 01-BFS that returns a distance array and a predecessor array.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 10 :: Int), (0, 2, 0), (2, 1, 1)]
-- >>> let gr = Gr.build 4 es
-- >>> let capacity = 10
-- >>> let (!dist, !prev) = Gr.trackingBfs01 4 (Gr.adjW gr) capacity (VU.singleton (0, 0))
-- >>> dist
-- [0,1,0,-1]
--
-- >>> Gr.constructPathFromRoot prev 1
-- [0,2,1]
--
-- @since 1.2.4.0
{-# INLINE trackingBfs01 #-}
trackingBfs01 ::
  forall i.
  (HasCallStack, Ix0 i, VU.Unbox i) =>
  -- | Zero-based index boundary.
  Bounds0 i ->
  -- | Graph function that takes the vertex and returns adjacent vertices with edge weights, where
  -- \(w > 0\).
  (i -> VU.Vector (i, Int)) ->
  -- | Capacity of deque, often the number of edges \(m\).
  Int ->
  -- | Weighted source vertices.
  VU.Vector (i, Int) ->
  -- | A tuple of (distance array in one-dimensional index, predecessor array). Unreachable vertices
  -- are assigned distance of @-1@.
  (VU.Vector Int, VU.Vector Int)
trackingBfs01 :: forall i.
(HasCallStack, Ix0 i, Unbox i) =>
i
-> (i -> Vector (i, Int))
-> Int
-> Vector (i, Int)
-> (Vector Int, Vector Int)
trackingBfs01 = Bool
-> i
-> (i -> Vector (i, Int))
-> Int
-> Vector (i, Int)
-> (Vector Int, Vector Int)
forall i.
(HasCallStack, Ix0 i, Unbox i) =>
Bool
-> i
-> (i -> Vector (i, Int))
-> Int
-> Vector (i, Int)
-> (Vector Int, Vector Int)
bfs01Impl Bool
True

{-# INLINEABLE bfs01Impl #-}
bfs01Impl ::
  forall i.
  (HasCallStack, Ix0 i, VU.Unbox i) =>
  Bool ->
  Bounds0 i ->
  (i -> VU.Vector (i, Int)) ->
  Int ->
  VU.Vector (i, Int) ->
  (VU.Vector Int, VU.Vector Int)
bfs01Impl :: forall i.
(HasCallStack, Ix0 i, Unbox i) =>
Bool
-> i
-> (i -> Vector (i, Int))
-> Int
-> Vector (i, Int)
-> (Vector Int, Vector Int)
bfs01Impl !Bool
trackPrev !i
bnd0 !i -> Vector (i, Int)
gr !Int
capacity !Vector (i, Int)
sources
  | Vector (i, Int) -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector (i, Int)
sources Bool -> Bool -> Bool
&& Bool
trackPrev = (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nVerts (-Int
1), Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nVerts (-Int
1))
  | Vector (i, Int) -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector (i, Int)
sources = (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nVerts (-Int
1), Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
0 (-Int
1))
  | Bool
otherwise = (forall s. ST s (Vector Int, Vector Int))
-> (Vector Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int, Vector Int))
 -> (Vector Int, Vector Int))
-> (forall s. ST s (Vector Int, Vector Int))
-> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
      MVector (PrimState (ST s)) Int
dist <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
nVerts Int
undef
      MVector (PrimState (ST s)) Int
prev <-
        if Bool
trackPrev
          then forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
nVerts (-Int
1)
          else forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
0 (-Int
1)
      -- NOTE: Just like Dijkstra, we need capacity of `m`, as the first appearance of a vertex is not
      -- always with minimum distance.
      Queue (PrimState (ST s)) (i, Int)
deque <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Queue (PrimState m) a)
Q.newDeque @_ @(i, Int) Int
capacity

      -- set source values
      Vector (i, Int) -> ((i, Int) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (i, Int)
sources (((i, Int) -> ST s ()) -> ST s ())
-> ((i, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!i
src, !Int
w0) -> do
        -- TODO: assert x1 <= w2
        let !i :: Int
i = i -> i -> Int
forall i. Ix0 i => i -> i -> Int
index0 i
bnd0 i
src
        !Int
lastD <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Int
dist Int
i
        -- Note that duplicate inputs are pruned here:
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lastD Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
undef) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
dist Int
i Int
w0
          Queue (PrimState (ST s)) (i, Int) -> (i, Int) -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
Q.pushBack Queue (PrimState (ST s)) (i, Int)
deque (i
src, Int
w0)

      let step :: i -> Int -> ST s ()
step !i
vExt0 !Int
w0 = do
            let !i0 :: Int
i0 = i -> i -> Int
forall i. Ix0 i => i -> i -> Int
index0 i
bnd0 i
vExt0
            !Int
wReserved0 <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Int
dist Int
i0
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
w0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
wReserved0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
              Vector (i, Int) -> ((i, Int) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (i -> Vector (i, Int)
gr i
vExt0) (((i, Int) -> ST s ()) -> ST s ())
-> ((i, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!i
vExt, !Int
dw) -> do
                let !w :: Int
w = Int
w0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dw
                let !i :: Int
i = i -> i -> Int
forall i. Ix0 i => i -> i -> Int
index0 i
bnd0 i
vExt
                !Int
wReserved <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Int
dist Int
i
                -- NOTE: Do pruning just like Dijkstra:
                Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
wReserved Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
undef Bool -> Bool -> Bool
|| Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
wReserved) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                  MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
dist Int
i Int
w
                  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
prev Int
i Int
i0
                  if Int
dw Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                    then Queue (PrimState (ST s)) (i, Int) -> (i, Int) -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
Q.pushFront Queue (PrimState (ST s)) (i, Int)
deque (i
vExt, Int
w)
                    else Queue (PrimState (ST s)) (i, Int) -> (i, Int) -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
Q.pushBack Queue (PrimState (ST s)) (i, Int)
deque (i
vExt, Int
w)

      (ST s () -> ST s ()) -> ST s ()
forall a. (a -> a) -> a
fix ((ST s () -> ST s ()) -> ST s ())
-> (ST s () -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ST s ()
popLoop -> do
        Queue (PrimState (ST s)) (i, Int) -> ST s (Maybe (i, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m (Maybe a)
Q.popFront Queue (PrimState (ST s)) (i, Int)
deque ST s (Maybe (i, Int)) -> (Maybe (i, Int) -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Maybe (i, Int)
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          Just (!i
vExt0, !Int
w0) -> do
            i -> Int -> ST s ()
step i
vExt0 Int
w0
            ST s ()
popLoop

      (,) (Vector Int -> Vector Int -> (Vector Int, Vector Int))
-> ST s (Vector Int)
-> ST s (Vector Int -> (Vector Int, Vector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState (ST s)) Int
dist ST s (Vector Int -> (Vector Int, Vector Int))
-> ST s (Vector Int) -> ST s (Vector Int, Vector Int)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState (ST s)) Int
prev
  where
    !undef :: Int
undef = -Int
1 :: Int
    !nVerts :: Int
nVerts = i -> Int
forall i. Ix0 i => i -> Int
rangeSize0 i
bnd0

-- | \(O((n + m) \log n)\) Dijkstra's algorithm that returns a distance array.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 10 :: Int), (1, 2, 20), (2, 3, 1), (1, 3, 40), (4, 3, 0)]
-- >>> let gr = Gr.build 5 es
-- >>> let capacity = 10
-- >>> Gr.dijkstra 5 (Gr.adjW gr) capacity (-1) (VU.singleton (0, 0))
-- [0,10,30,31,-1]
--
-- @since 1.2.4.0
{-# INLINE dijkstra #-}
dijkstra ::
  forall i w.
  (HasCallStack, Ix0 i, Ord i, VU.Unbox i, Num w, Ord w, VU.Unbox w) =>
  -- | Zero-based vertex boundary.
  Bounds0 i ->
  -- | Graph function that takes a vertex and returns adjacent vertices with edge weights, where
  -- \(w \ge 0\).
  (i -> VU.Vector (i, w)) ->
  -- | Capacity of the heap, often the number of edges \(m\).
  Int ->
  -- | Distance assignment for unreachable vertices.
  w ->
  -- | Source vertices with initial weights.
  VU.Vector (i, w) ->
  -- | Distance array in one-dimensional index.
  VU.Vector w
dijkstra :: forall i w.
(HasCallStack, Ix0 i, Ord i, Unbox i, Num w, Ord w, Unbox w) =>
i -> (i -> Vector (i, w)) -> Int -> w -> Vector (i, w) -> Vector w
dijkstra !i
bnd0 !i -> Vector (i, w)
gr !Int
capacity !w
undefW !Vector (i, w)
sources =
  let (!Vector w
dist, !Vector Int
_) = Bool
-> i
-> (i -> Vector (i, w))
-> Int
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
forall i w.
(HasCallStack, Ix0 i, Ord i, Unbox i, Num w, Ord w, Unbox w) =>
Bool
-> i
-> (i -> Vector (i, w))
-> Int
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
dijkstraImpl Bool
False i
bnd0 i -> Vector (i, w)
gr Int
capacity w
undefW Vector (i, w)
sources
   in Vector w
dist

-- | \(O((n + m) \log n)\) Dijkstra's algorithm that returns a distance array and a predecessor
-- array.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 10 :: Int), (1, 2, 20), (2, 3, 1), (1, 3, 40), (4, 3, 0)]
-- >>> let gr = Gr.build 5 es
-- >>> let capacity = 10
-- >>> let (!dist, !prev) = Gr.trackingDijkstra 5 (Gr.adjW gr) capacity (-1) (VU.singleton (0, 0))
-- >>> dist
-- [0,10,30,31,-1]
--
-- >>> Gr.constructPathFromRoot prev 3
-- [0,1,2,3]
--
-- @since 1.2.4.0
{-# INLINE trackingDijkstra #-}
trackingDijkstra ::
  forall i w.
  (HasCallStack, Ix0 i, Ord i, VU.Unbox i, Num w, Ord w, VU.Unbox w) =>
  -- | Zero-based vertex boundary.
  Bounds0 i ->
  -- | Graph function that takes a vertex and returns adjacent vertices with edge weights, where
  -- \(w \ge 0\).
  (i -> VU.Vector (i, w)) ->
  -- | Capacity of the heap, often the number of edges \(m\).
  Int ->
  -- | Distance assignment for unreachable vertices.
  w ->
  -- | Source vertices with weights.
  VU.Vector (i, w) ->
  -- | A tuple of (distance array in one-dimensional index, predecessor array).
  (VU.Vector w, VU.Vector Int)
trackingDijkstra :: forall i w.
(HasCallStack, Ix0 i, Ord i, Unbox i, Num w, Ord w, Unbox w) =>
i
-> (i -> Vector (i, w))
-> Int
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
trackingDijkstra = Bool
-> i
-> (i -> Vector (i, w))
-> Int
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
forall i w.
(HasCallStack, Ix0 i, Ord i, Unbox i, Num w, Ord w, Unbox w) =>
Bool
-> i
-> (i -> Vector (i, w))
-> Int
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
dijkstraImpl Bool
True

{-# INLINEABLE dijkstraImpl #-}
dijkstraImpl ::
  forall i w.
  (HasCallStack, Ix0 i, Ord i, VU.Unbox i, Num w, Ord w, VU.Unbox w) =>
  Bool ->
  Bounds0 i ->
  (i -> VU.Vector (i, w)) ->
  Int ->
  w ->
  VU.Vector (i, w) ->
  (VU.Vector w, VU.Vector Int)
dijkstraImpl :: forall i w.
(HasCallStack, Ix0 i, Ord i, Unbox i, Num w, Ord w, Unbox w) =>
Bool
-> i
-> (i -> Vector (i, w))
-> Int
-> w
-> Vector (i, w)
-> (Vector w, Vector Int)
dijkstraImpl !Bool
trackPrev !i
bnd0 !i -> Vector (i, w)
gr !Int
capacity !w
undefW !Vector (i, w)
sources
  | Vector (i, w) -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector (i, w)
sources Bool -> Bool -> Bool
&& Bool
trackPrev = (Int -> w -> Vector w
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nVerts w
undefW, Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nVerts (-Int
1))
  | Vector (i, w) -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector (i, w)
sources = (Int -> w -> Vector w
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nVerts w
undefW, Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
0 (-Int
1))
  | Bool
otherwise = (forall s. ST s (Vector w, Vector Int)) -> (Vector w, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector w, Vector Int)) -> (Vector w, Vector Int))
-> (forall s. ST s (Vector w, Vector Int))
-> (Vector w, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
      !MVector s w
dist <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @w Int
nVerts w
undefW
      -- REMARK: (w, i) for sort by width
      !Heap s (w, i)
heap <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Heap (PrimState m) a)
MH.new @_ @(w, i) Int
capacity
      !MVector s Int
prev <-
        if Bool
trackPrev
          then forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
nVerts (-Int
1)
          else forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
0 (-Int
1)

      Vector (i, w) -> ((i, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (i, w)
sources (((i, w) -> ST s ()) -> ST s ()) -> ((i, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!i
v, !w
w) -> do
        let !i :: Int
i = i -> i -> Int
forall i. Ix0 i => i -> i -> Int
index0 i
bnd0 i
v
        MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState (ST s)) w
dist Int
i w
w
        Heap (PrimState (ST s)) (w, i) -> (w, i) -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Ord a, Unbox a) =>
Heap (PrimState m) a -> a -> m ()
MH.push Heap s (w, i)
Heap (PrimState (ST s)) (w, i)
heap (w
w, i
v)

      (ST s () -> ST s ()) -> ST s ()
forall a. (a -> a) -> a
fix ((ST s () -> ST s ()) -> ST s ())
-> (ST s () -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ST s ()
loop -> do
        Heap (PrimState (ST s)) (w, i) -> ST s (Maybe (w, i))
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Ord a, Unbox a) =>
Heap (PrimState m) a -> m (Maybe a)
MH.pop Heap s (w, i)
Heap (PrimState (ST s)) (w, i)
heap ST s (Maybe (w, i)) -> (Maybe (w, i) -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Maybe (w, i)
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          Just (!w
w1, !i
v1) -> do
            let !i1 :: Int
i1 = i -> i -> Int
forall i. Ix0 i => i -> i -> Int
index0 i
bnd0 i
v1
            !w
wReserved <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
dist Int
i1
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
wReserved w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
w1) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
              Vector (i, w) -> ((i, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (i -> Vector (i, w)
gr i
v1) (((i, w) -> ST s ()) -> ST s ()) -> ((i, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!i
v2, !w
dw2) -> do
                let !i2 :: Int
i2 = i -> i -> Int
forall i. Ix0 i => i -> i -> Int
index0 i
bnd0 i
v2
                !w
w2 <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
dist Int
i2
                let !w2' :: w
w2' = w
w1 w -> w -> w
forall a. Num a => a -> a -> a
+ w
dw2
                Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
w2 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW Bool -> Bool -> Bool
|| w
w2' w -> w -> Bool
forall a. Ord a => a -> a -> Bool
< w
w2) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                  MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState (ST s)) w
dist Int
i2 w
w2'
                  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
prev Int
i2 Int
i1
                  Heap (PrimState (ST s)) (w, i) -> (w, i) -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Ord a, Unbox a) =>
Heap (PrimState m) a -> a -> m ()
MH.push Heap s (w, i)
Heap (PrimState (ST s)) (w, i)
heap (w
w2', i
v2)
            ST s ()
loop

      (,) (Vector w -> Vector Int -> (Vector w, Vector Int))
-> ST s (Vector w) -> ST s (Vector Int -> (Vector w, Vector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) w -> ST s (Vector w)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s w
MVector (PrimState (ST s)) w
dist ST s (Vector Int -> (Vector w, Vector Int))
-> ST s (Vector Int) -> ST s (Vector w, Vector Int)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
prev
  where
    !nVerts :: Int
nVerts = i -> Int
forall i. Ix0 i => i -> Int
rangeSize0 i
bnd0

-- -- | Option for `bellmanFord`.
-- data BellmanFordPolicy = QuitOnNegaitveLoop | ContinueOnNegaitveLoop

-- | \(O(nm)\) Bellman–ford algorithm that returns a distance array, or `Nothing` on negative loop
-- detection. Vertices are one-dimensional.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let gr = Gr.build @Int 5 $ VU.fromList [(0, 1, 10), (1, 2, -20), (2, 3, 1), (1, 3, 40), (4, 3, 0)]
-- >>> let undefW = maxBound `div` 2
-- >>> Gr.bellmanFord 5 (Gr.adjW gr) undefW (VU.singleton (0, 0))
-- Just [0,10,-10,-9,4611686018427387903]
--
-- It returns `Nothing` on negative loop detection:
--
-- >>> let gr = Gr.build @Int 2 $ VU.fromList [(0, 1, -1), (1, 0, -1)]
-- >>> Gr.bellmanFord 5 (Gr.adjW gr) undefW (VU.singleton (0, 0))
-- Nothing
--
-- @since 1.2.4.0
{-# INLINE bellmanFord #-}
bellmanFord ::
  forall w.
  (HasCallStack, Num w, Ord w, VU.Unbox w) =>
  -- | The number of vertices.
  Int ->
  -- | Graph function. Edges weights can be negative.
  (Int -> VU.Vector (Int, w)) ->
  -- | Distance assignment for unreachable vertices.
  w ->
  -- | Source vertex with initial distances.
  VU.Vector (Int, w) ->
  -- | Distance array in one-dimensional index.
  Maybe (VU.Vector w)
bellmanFord :: forall w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Int
-> (Int -> Vector (Int, w))
-> w
-> Vector (Int, w)
-> Maybe (Vector w)
bellmanFord {- !policy -} !Int
nVerts !Int -> Vector (Int, w)
gr !w
undefW Vector (Int, w)
source = do
  (!Vector w
dist, !Vector Int
_) <- Bool
-> Int
-> (Int -> Vector (Int, w))
-> w
-> Vector (Int, w)
-> Maybe (Vector w, Vector Int)
forall w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> Int
-> (Int -> Vector (Int, w))
-> w
-> Vector (Int, w)
-> Maybe (Vector w, Vector Int)
bellmanFordImpl Bool
False Int
nVerts Int -> Vector (Int, w)
gr w
undefW Vector (Int, w)
source
  Vector w -> Maybe (Vector w)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector w
dist

-- | \(O(nm)\) Bellman–ford algorithm that returns a distance array and a predecessor array, or
-- `Nothing` on negative loop detection. Vertices are one-dimensional.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let gr = Gr.build @Int 5 $ VU.fromList [(0, 1, 10), (1, 2, -20), (2, 3, 1), (1, 3, 40), (4, 3, 0)]
-- >>> let undefW = maxBound `div` 2
-- >>> let Just (!dist, !prev) = Gr.trackingBellmanFord 5 (Gr.adjW gr) undefW (VU.singleton (0, 0))
-- >>> dist
-- [0,10,-10,-9,4611686018427387903]
--
-- >>> Gr.constructPathFromRoot prev 3
-- [0,1,2,3]
--
-- It returns `Nothing` on negative loop detection:
--
-- >>> let gr = Gr.build @Int 2 $ VU.fromList [(0, 1, -1), (1, 0, -1)]
-- >>> Gr.trackingBellmanFord 5 (Gr.adjW gr) undefW (VU.singleton (0, 0))
-- Nothing
--
-- @since 1.2.4.0
{-# INLINE trackingBellmanFord #-}
trackingBellmanFord ::
  forall w.
  (HasCallStack, Num w, Ord w, VU.Unbox w) =>
  -- | The number of vertices.
  Int ->
  -- | Graph function. The weight can be negative.
  (Int -> VU.Vector (Int, w)) ->
  -- | Distance assignment for unreachable vertices.
  w ->
  -- | Source vertex with initial distances.
  VU.Vector (Int, w) ->
  -- | A tuple of (distance array, predecessor array).
  Maybe (VU.Vector w, VU.Vector Int)
trackingBellmanFord :: forall w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Int
-> (Int -> Vector (Int, w))
-> w
-> Vector (Int, w)
-> Maybe (Vector w, Vector Int)
trackingBellmanFord {- !policy -} = Bool
-> Int
-> (Int -> Vector (Int, w))
-> w
-> Vector (Int, w)
-> Maybe (Vector w, Vector Int)
forall w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> Int
-> (Int -> Vector (Int, w))
-> w
-> Vector (Int, w)
-> Maybe (Vector w, Vector Int)
bellmanFordImpl Bool
True

{-# INLINEABLE bellmanFordImpl #-}
bellmanFordImpl ::
  forall w.
  (HasCallStack, Num w, Ord w, VU.Unbox w) =>
  Bool ->
  Int ->
  (Int -> VU.Vector (Int, w)) ->
  w ->
  VU.Vector (Int, w) ->
  Maybe (VU.Vector w, VU.Vector Int)
bellmanFordImpl :: forall w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> Int
-> (Int -> Vector (Int, w))
-> w
-> Vector (Int, w)
-> Maybe (Vector w, Vector Int)
bellmanFordImpl {- !policy -} !Bool
trackPrev !Int
nVerts !Int -> Vector (Int, w)
gr !w
undefW !Vector (Int, w)
sources = (forall s. ST s (Maybe (Vector w, Vector Int)))
-> Maybe (Vector w, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe (Vector w, Vector Int)))
 -> Maybe (Vector w, Vector Int))
-> (forall s. ST s (Maybe (Vector w, Vector Int)))
-> Maybe (Vector w, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  !MVector (PrimState (ST s)) w
dist <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @w Int
nVerts w
undefW
  !MVector (PrimState (ST s)) Int
prev <-
    if Bool
trackPrev
      then forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
nVerts (-Int
1)
      else forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
0 (-Int
1)

  Vector (Int, w) -> ((Int, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (Int, w)
sources (((Int, w) -> ST s ()) -> ST s ())
-> ((Int, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
v, !w
w) -> do
    !w
lastD <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) w
dist Int
v
    -- Note that duplicate inputs are pruned here:
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
lastD w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) w
dist Int
v w
w
  MVector (PrimState (ST s)) Bool
updated <- Int -> Bool -> ST s (MVector (PrimState (ST s)) Bool)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 Bool
False

  -- look around adjaenct vertices
  let update :: Int -> ST s ()
update Int
v1 = do
        w
d1 <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) w
dist Int
v1
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
d1 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
/= w
undefW) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          Vector (Int, w) -> ((Int, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Int -> Vector (Int, w)
gr Int
v1) (((Int, w) -> ST s ()) -> ST s ())
-> ((Int, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
v2, !w
dw) -> do
            w
d2 <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) w
dist Int
v2
            let !d2' :: w
d2' = w
d1 w -> w -> w
forall a. Num a => a -> a -> a
+ w
dw
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
d2 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW Bool -> Bool -> Bool
|| w
d2' w -> w -> Bool
forall a. Ord a => a -> a -> Bool
< w
d2) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
              MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) w
dist Int
v2 w
d2'
              Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
prev Int
v2 Int
v1
              -- NOTE: we should actually instantly stop if nLoop == nVerts + 1, but
              -- here we're preferring simple code. Be warned that we're not correctly handling
              -- the distance array on negative loop.
              MVector (PrimState (ST s)) Bool -> Int -> Bool -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Bool
updated Int
0 Bool
True

  let runLoop :: Int -> ST s (Maybe (Vector w, Vector Int))
runLoop Int
nLoop
        | Int
nLoop Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 = do
            -- We detected update in the (n + 1)-th loop, so we found negative loop
            Maybe (Vector w, Vector Int) -> ST s (Maybe (Vector w, Vector Int))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Vector w, Vector Int)
forall a. Maybe a
Nothing
        | Bool
otherwise = do
            [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] Int -> ST s ()
update
            Bool
b <- MVector (PrimState (ST s)) Bool -> Int -> Bool -> ST s Bool
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector (PrimState (ST s)) Bool
updated Int
0 Bool
False
            if Bool
b
              then Int -> ST s (Maybe (Vector w, Vector Int))
runLoop (Int
nLoop Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
              else (Vector w, Vector Int) -> Maybe (Vector w, Vector Int)
forall a. a -> Maybe a
Just ((Vector w, Vector Int) -> Maybe (Vector w, Vector Int))
-> ST s (Vector w, Vector Int)
-> ST s (Maybe (Vector w, Vector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((,) (Vector w -> Vector Int -> (Vector w, Vector Int))
-> ST s (Vector w) -> ST s (Vector Int -> (Vector w, Vector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) w -> ST s (Vector w)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState (ST s)) w
dist ST s (Vector Int -> (Vector w, Vector Int))
-> ST s (Vector Int) -> ST s (Vector w, Vector Int)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState (ST s)) Int
prev)

  Int -> ST s (Maybe (Vector w, Vector Int))
runLoop Int
0

-- | \(O(n^3)\) Floyd–Warshall algorithm that returns a distance matrix \(m\), which should be
-- accessed as @m VU.! (`index0` (n, n) (from, to))@. Negative loop can be detected by testing if
-- there's any vertex \(v\) where @m VU.! (`index0` (n, n) (v, v))@.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 10 :: Int), (1, 2, -20), (2, 3, 1), (1, 3, 40), (4, 3, 0)]
-- >>> let undefW = maxBound `div` 2
-- >>> let dist = Gr.floydWarshall 5 es undefW
-- >>> dist VG.! (5 * 0 + 3) -- from `0` to `3`
-- -9
--
-- >>> dist VG.! (5 * 1 + 3) -- from `0` to `3`
-- -19
--
-- Negative loop can be detected by testing if there's any vertex \(v\) where
-- @m VU.! (`index0` (n, n) (v, v))@:
--
-- >>> any (\v -> dist VG.! (5 * v + v) < 0) [0 .. 5 - 1]
-- False
--
-- >>> let es = VU.fromList [(0, 1, -1 :: Int), (1, 0, -1)]
-- >>> let dist = Gr.floydWarshall 3 es undefW
-- >>> any (\v -> dist VG.! (3 * v + v) < 0) [0 .. 3 - 1]
-- True
--
-- @since 1.2.4.0
{-# INLINE floydWarshall #-}
floydWarshall ::
  forall w.
  (HasCallStack, Num w, Ord w, VU.Unbox w) =>
  -- | The number of vertices.
  Int ->
  -- | Weighted edges.
  VU.Vector (Int, Int, w) ->
  -- | Distance assignment \(d_0 \gt 0\) for unreachable vertices. It should be @maxBound `div` 2@
  -- for `Int`.
  w ->
  -- | Distance array in one-dimensional index.
  VU.Vector w
floydWarshall :: forall w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Int -> Vector (Int, Int, w) -> w -> Vector w
floydWarshall !Int
nVerts !Vector (Int, Int, w)
edges !w
undefW = (forall s. ST s (MVector s w)) -> Vector w
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s w)) -> Vector w)
-> (forall s. ST s (MVector s w)) -> Vector w
forall a b. (a -> b) -> a -> b
$ do
  (!MVector s w
dist, !MVector s Int
_) <- Bool
-> Int
-> Vector (Int, Int, w)
-> w
-> ST s (MVector s w, MVector s Int)
forall s w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> Int
-> Vector (Int, Int, w)
-> w
-> ST s (MVector s w, MVector s Int)
newFloydWarshallST Bool
False Int
nVerts Vector (Int, Int, w)
edges w
undefW
  MVector s w -> ST s (MVector s w)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s w
dist

-- | \(O(n^3)\) Floyd–Warshall algorithm that returns a distance matrix \(m\) and predecessor
-- matrix \(p\). The distance matrix should be accessed as @m VU.! (`index0` (n, n) (from, to))@,
-- and the predecessor matrix should be accessed as @m VU.! (`index0` (n, n) (root, v))@. There's a
-- negative loop if there's any vertex \(v\) where @m VU.! (`index0` (n, n) (v, v))@.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 10 :: Int), (1, 2, -20), (2, 3, 1), (1, 3, 40), (4, 3, 0)]
-- >>> let undefW = maxBound `div` 2
-- >>> let (!dist, !prev) = Gr.trackingFloydWarshall 5 es undefW
-- >>> dist VG.! (5 * 0 + 3) -- from `0` to `3`
-- -9
--
-- >>> Gr.constructPathFromRootMat prev 0 3 -- from `0` to `3`
-- [0,1,2,3]
--
-- >>> dist VG.! (5 * 1 + 3) -- from `0` to `3`
-- -19
--
-- >>> Gr.constructPathFromRootMat prev 1 3 -- from `1` to `3`
-- [1,2,3]
--
-- Negative loop can be detected by testing if there's any vertex \(v\) where
-- @m VU.! (`index0` (n, n) (v, v))@:
--
-- >>> any (\v -> dist VG.! (5 * v + v) < 0) [0 .. 5 - 1]
-- False
--
-- >>> let es = VU.fromList [(0, 1, -1 :: Int), (1, 0, -1)]
-- >>> let (!dist, !_) = Gr.trackingFloydWarshall 3 es undefW
-- >>> any (\v -> dist VG.! (3 * v + v) < 0) [0 .. 3 - 1]
-- True
--
-- @since 1.2.4.0
{-# INLINE trackingFloydWarshall #-}
trackingFloydWarshall ::
  forall w.
  (HasCallStack, Num w, Ord w, VU.Unbox w) =>
  -- | The number of vertices.
  Int ->
  -- | Weighted edges.
  VU.Vector (Int, Int, w) ->
  -- | Distance assignment \(d_0 \gt 0\) for unreachable vertices. It should be @maxBound `div` 2@
  -- for `Int`.
  w ->
  -- | Distance array in one-dimensional index.
  (VU.Vector w, VU.Vector Int)
trackingFloydWarshall :: forall w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Int -> Vector (Int, Int, w) -> w -> (Vector w, Vector Int)
trackingFloydWarshall !Int
nVerts !Vector (Int, Int, w)
edges !w
undefW = (forall s. ST s (Vector w, Vector Int)) -> (Vector w, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector w, Vector Int)) -> (Vector w, Vector Int))
-> (forall s. ST s (Vector w, Vector Int))
-> (Vector w, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  (!MVector s w
dist, !MVector s Int
prev) <- Bool
-> Int
-> Vector (Int, Int, w)
-> w
-> ST s (MVector s w, MVector s Int)
forall s w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> Int
-> Vector (Int, Int, w)
-> w
-> ST s (MVector s w, MVector s Int)
newFloydWarshallST Bool
True Int
nVerts Vector (Int, Int, w)
edges w
undefW
  (,) (Vector w -> Vector Int -> (Vector w, Vector Int))
-> ST s (Vector w) -> ST s (Vector Int -> (Vector w, Vector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) w -> ST s (Vector w)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s w
MVector (PrimState (ST s)) w
dist ST s (Vector Int -> (Vector w, Vector Int))
-> ST s (Vector Int) -> ST s (Vector w, Vector Int)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
prev

-- | \(O(n^3)\) Floyd–Warshall algorithm that returns a distance matrix \(m\), which should be
-- accessed as @m VU.! (n * from + to)@. There's a negative cycle if any @m VU.! (n * i + i)@ is
-- negative.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 1 :: Int), (1, 2, 1), (2, 3, 1), (1, 3, 4)]
-- >>> let undefW = -1
-- >>> dist <- Gr.newFloydWarshall 4 es undefW
-- >>> VGM.read dist (4 * 0 + 3)
-- 3
--
-- >>> updateEdgeFloydWarshall dist 4 undefW 1 3 (-2)
-- >>> VGM.read dist (4 * 0 + 3)
-- -1
--
-- @since 1.2.4.0
{-# INLINE newFloydWarshall #-}
newFloydWarshall ::
  forall m w.
  (HasCallStack, PrimMonad m, Num w, Ord w, VU.Unbox w) =>
  -- | The number of vertices.
  Int ->
  -- | Weighted edges.
  VU.Vector (Int, Int, w) ->
  -- | Distance assignment for unreachable vertices.
  w ->
  -- | Distance array in one-dimensional index.
  m (VUM.MVector (PrimState m) w)
newFloydWarshall :: forall (m :: * -> *) w.
(HasCallStack, PrimMonad m, Num w, Ord w, Unbox w) =>
Int -> Vector (Int, Int, w) -> w -> m (MVector (PrimState m) w)
newFloydWarshall !Int
nVerts !Vector (Int, Int, w)
edges !w
undefW = ST (PrimState m) (MVector (PrimState m) w)
-> m (MVector (PrimState m) w)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (MVector (PrimState m) w)
 -> m (MVector (PrimState m) w))
-> ST (PrimState m) (MVector (PrimState m) w)
-> m (MVector (PrimState m) w)
forall a b. (a -> b) -> a -> b
$ do
  (!MVector (PrimState m) w
dist, !MVector (PrimState m) Int
_) <- Bool
-> Int
-> Vector (Int, Int, w)
-> w
-> ST
     (PrimState m) (MVector (PrimState m) w, MVector (PrimState m) Int)
forall s w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> Int
-> Vector (Int, Int, w)
-> w
-> ST s (MVector s w, MVector s Int)
newFloydWarshallST Bool
False Int
nVerts Vector (Int, Int, w)
edges w
undefW
  MVector (PrimState m) w
-> ST (PrimState m) (MVector (PrimState m) w)
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector (PrimState m) w
dist

-- | \(O(n^3)\) Floyd–Warshall algorithm that returns a distance matrix \(m\) and predecessor
-- matrix. There's a negative cycle if any @m VU.! (n * i + i)@ is negative.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let es = VU.fromList [(0, 1, 1 :: Int), (1, 2, 1), (2, 3, 1), (1, 3, 4)]
-- >>> let undefW = -1
-- >>> (!dist, !prev) <- Gr.newTrackingFloydWarshall 4 es undefW
-- >>> VGM.read dist (4 * 0 + 3)
-- 3
--
-- >>> constructPathFromRootMatM prev 0 3
-- [0,1,2,3]
--
-- >>> updateEdgeTrackingFloydWarshall dist prev 4 undefW 1 3 (-2)
-- >>> VGM.read dist (4 * 0 + 3)
-- -1
--
-- >>> constructPathFromRootMatM prev 0 3
-- [0,1,3]
--
-- @since 1.2.4.0
{-# INLINE newTrackingFloydWarshall #-}
newTrackingFloydWarshall ::
  forall m w.
  (HasCallStack, PrimMonad m, Num w, Ord w, VU.Unbox w) =>
  -- | The number of vertices.
  Int ->
  -- | Weighted edges.
  VU.Vector (Int, Int, w) ->
  -- | Distance assignment for unreachable vertices.
  w ->
  -- | Distance array in one-dimensional index.
  m (VUM.MVector (PrimState m) w, VUM.MVector (PrimState m) Int)
newTrackingFloydWarshall :: forall (m :: * -> *) w.
(HasCallStack, PrimMonad m, Num w, Ord w, Unbox w) =>
Int
-> Vector (Int, Int, w)
-> w
-> m (MVector (PrimState m) w, MVector (PrimState m) Int)
newTrackingFloydWarshall !Int
nVerts !Vector (Int, Int, w)
edges !w
undefW = ST
  (PrimState m) (MVector (PrimState m) w, MVector (PrimState m) Int)
-> m (MVector (PrimState m) w, MVector (PrimState m) Int)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST
   (PrimState m) (MVector (PrimState m) w, MVector (PrimState m) Int)
 -> m (MVector (PrimState m) w, MVector (PrimState m) Int))
-> ST
     (PrimState m) (MVector (PrimState m) w, MVector (PrimState m) Int)
-> m (MVector (PrimState m) w, MVector (PrimState m) Int)
forall a b. (a -> b) -> a -> b
$ do
  Bool
-> Int
-> Vector (Int, Int, w)
-> w
-> ST
     (PrimState m) (MVector (PrimState m) w, MVector (PrimState m) Int)
forall s w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> Int
-> Vector (Int, Int, w)
-> w
-> ST s (MVector s w, MVector s Int)
newFloydWarshallST Bool
True Int
nVerts Vector (Int, Int, w)
edges w
undefW

{-# INLINEABLE newFloydWarshallST #-}
newFloydWarshallST ::
  forall s w.
  (HasCallStack, Num w, Ord w, VU.Unbox w) =>
  Bool ->
  Int ->
  VU.Vector (Int, Int, w) ->
  w ->
  ST s (VUM.MVector s w, VUM.MVector s Int)
newFloydWarshallST :: forall s w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> Int
-> Vector (Int, Int, w)
-> w
-> ST s (MVector s w, MVector s Int)
newFloydWarshallST !Bool
trackPrev !Int
nVerts !Vector (Int, Int, w)
edges !w
undefW = do
  !MVector s w
dist <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @w (Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nVerts) w
undefW
  !MVector s Int
prev <-
    if Bool
trackPrev
      then forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int (Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nVerts) (-Int
1)
      else forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
0 (-Int
1)

  -- diagonals (self to self)
  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v -> do
    MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState (ST s)) w
dist (Int -> Int -> Int
idx Int
v Int
v) w
0

  -- initial walks
  Vector (Int, Int, w) -> ((Int, Int, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (Int, Int, w)
edges (((Int, Int, w) -> ST s ()) -> ST s ())
-> ((Int, Int, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
v1, !Int
v2, !w
dw) -> do
    let !i :: Int
i = Int -> Int -> Int
idx Int
v1 Int
v2
    w
wOld <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
dist Int
i
    -- REMARK: We're handling multiple edges here:
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
wOld w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW Bool -> Bool -> Bool
|| w
dw w -> w -> Bool
forall a. Ord a => a -> a -> Bool
< w
wOld) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState (ST s)) w
dist Int
i w
dw
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
prev Int
i Int
v1

  -- N times update
  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
via -> do
    -- update
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
from -> do
      [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
to -> do
        let !iFromTo :: Int
iFromTo = Int -> Int -> Int
idx Int
from Int
to
        !w
w1 <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
dist Int
iFromTo
        !w
w2 <- do
          !w
d1 <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
dist (Int -> ST s w) -> Int -> ST s w
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
idx Int
from Int
via
          !w
d2 <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
dist (Int -> ST s w) -> Int -> ST s w
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
idx Int
via Int
to
          w -> ST s w
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (w -> ST s w) -> w -> ST s w
forall a b. (a -> b) -> a -> b
$! if w
d1 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW Bool -> Bool -> Bool
|| w
d2 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW then w
undefW else w
d1 w -> w -> w
forall a. Num a => a -> a -> a
+ w
d2
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
w2 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
/= w
undefW Bool -> Bool -> Bool
&& (w
w1 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW Bool -> Bool -> Bool
|| w
w2 w -> w -> Bool
forall a. Ord a => a -> a -> Bool
< w
w1)) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState (ST s)) w
dist Int
iFromTo w
w2
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
prev Int
iFromTo (Int -> ST s ()) -> ST s Int -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
prev (Int -> Int -> Int
idx Int
via Int
to)

  (MVector s w, MVector s Int) -> ST s (MVector s w, MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MVector s w
dist, MVector s Int
prev)
  where
    idx :: Int -> Int -> Int
idx !Int
from !Int
to = Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
from Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
to

-- | \(O(n^2)\) Updates distance matrix of Floyd–Warshall on edge weight decreasement or new edge
-- addition.
--
-- @since 1.2.4.0
{-# INLINE updateEdgeFloydWarshall #-}
updateEdgeFloydWarshall ::
  forall m w.
  (HasCallStack, PrimMonad m, Num w, Ord w, VU.Unbox w) =>
  -- | Distance matrix.
  VUM.MVector (PrimState m) w ->
  -- | The number of vertices.
  Int ->
  -- | Distance assignment \(d_0 \gt 0\) for unreachable vertices. It should be @maxBound `div` 2@
  -- for `Int`.
  w ->
  -- | Edge information: @from@ vertex.
  Int ->
  -- | Edge information: @to@ vertex.
  Int ->
  -- | Edge information: @weight@ vertex.
  w ->
  -- | Distance array in one-dimensional index.
  m ()
updateEdgeFloydWarshall :: forall (m :: * -> *) w.
(HasCallStack, PrimMonad m, Num w, Ord w, Unbox w) =>
MVector (PrimState m) w -> Int -> w -> Int -> Int -> w -> m ()
updateEdgeFloydWarshall MVector (PrimState m) w
mat Int
nVerts w
undefW Int
a Int
b w
w = do
  MVector (PrimState m) Int
prev <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Int Int
0 (-Int
1 :: Int)
  ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ Bool
-> MVector (PrimState m) w
-> MVector (PrimState m) Int
-> Int
-> w
-> Int
-> Int
-> w
-> ST (PrimState m) ()
forall s w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> MVector s w
-> MVector s Int
-> Int
-> w
-> Int
-> Int
-> w
-> ST s ()
updateEdgeFloydWarshallST Bool
False MVector (PrimState m) w
mat MVector (PrimState m) Int
prev Int
nVerts w
undefW Int
a Int
b w
w

-- | \(O(n^2)\) Updates distance matrix of Floyd–Warshall on edge weight decreasement or new edge
-- addition.
--
-- @since 1.2.4.0
{-# INLINE updateEdgeTrackingFloydWarshall #-}
updateEdgeTrackingFloydWarshall ::
  forall m w.
  (HasCallStack, PrimMonad m, Num w, Ord w, VU.Unbox w) =>
  -- | Distance matrix.
  VUM.MVector (PrimState m) w ->
  -- | Predecessor matrix.
  VUM.MVector (PrimState m) Int ->
  -- | The number of vertices.
  Int ->
  -- | Distance assignment \(d_0 \gt 0\) for unreachable vertices. It should be @maxBound `div` 2@
  -- for `Int`.
  w ->
  -- | Edge information: @from@ vertex.
  Int ->
  -- | Edge information: @to@ vertex.
  Int ->
  -- | Edge information: @weight@ vertex.
  w ->
  -- | Distance array in one-dimensional index.
  m ()
updateEdgeTrackingFloydWarshall :: forall (m :: * -> *) w.
(HasCallStack, PrimMonad m, Num w, Ord w, Unbox w) =>
MVector (PrimState m) w
-> MVector (PrimState m) Int -> Int -> w -> Int -> Int -> w -> m ()
updateEdgeTrackingFloydWarshall MVector (PrimState m) w
mat MVector (PrimState m) Int
prev Int
nVerts w
undefW Int
a Int
b w
w = do
  ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ Bool
-> MVector (PrimState m) w
-> MVector (PrimState m) Int
-> Int
-> w
-> Int
-> Int
-> w
-> ST (PrimState m) ()
forall s w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> MVector s w
-> MVector s Int
-> Int
-> w
-> Int
-> Int
-> w
-> ST s ()
updateEdgeFloydWarshallST Bool
True MVector (PrimState m) w
mat MVector (PrimState m) Int
prev Int
nVerts w
undefW Int
a Int
b w
w

-- O(2) update floyd warshall on edge weight decreasement or edge addition
-- https://www.slideshare.net/chokudai/arc035 - C
{-# INLINEABLE updateEdgeFloydWarshallST #-}
updateEdgeFloydWarshallST ::
  forall s w.
  (HasCallStack, Num w, Ord w, VU.Unbox w) =>
  Bool ->
  VUM.MVector s w ->
  VUM.MVector s Int ->
  Int ->
  w ->
  Int ->
  Int ->
  w ->
  ST s ()
updateEdgeFloydWarshallST :: forall s w.
(HasCallStack, Num w, Ord w, Unbox w) =>
Bool
-> MVector s w
-> MVector s Int
-> Int
-> w
-> Int
-> Int
-> w
-> ST s ()
updateEdgeFloydWarshallST Bool
trackPrev MVector s w
mat MVector s Int
prev Int
nVerts w
undefW Int
a Int
b w
dw = do
  w
wOld0 <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
mat (Int -> ST s w) -> Int -> ST s w
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
idx Int
a Int
b
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
wOld0 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW Bool -> Bool -> Bool
|| w
dw w -> w -> Bool
forall a. Ord a => a -> a -> Bool
< w
wOld0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState (ST s)) w
mat (Int -> Int -> Int
idx Int
a Int
b) w
dw
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
prev (Int -> Int -> Int
idx Int
a Int
b) Int
a
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
from -> do
      [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
to -> do
        w
wOld <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
mat (Int -> ST s w) -> Int -> ST s w
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
idx Int
from Int
to

        w
w' <- do
          w
ia <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
mat (Int -> ST s w) -> Int -> ST s w
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
idx Int
from Int
a
          w
bj <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
mat (Int -> ST s w) -> Int -> ST s w
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
idx Int
b Int
to
          let w1 :: w
w1
                | w
ia w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW Bool -> Bool -> Bool
|| w
bj w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW = w
undefW
                | Bool
otherwise = w
ia w -> w -> w
forall a. Num a => a -> a -> a
+ w
dw w -> w -> w
forall a. Num a => a -> a -> a
+ w
bj

          w
ib <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
mat (Int -> ST s w) -> Int -> ST s w
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
idx Int
from Int
b
          w
aj <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s w
MVector (PrimState (ST s)) w
mat (Int -> ST s w) -> Int -> ST s w
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
idx Int
a Int
to
          let w2 :: w
w2
                | w
ib w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW Bool -> Bool -> Bool
|| w
aj w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW = w
undefW
                | Bool
otherwise = w
ib w -> w -> w
forall a. Num a => a -> a -> a
+ w
dw w -> w -> w
forall a. Num a => a -> a -> a
+ w
aj

          w -> ST s w
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (w -> ST s w) -> w -> ST s w
forall a b. (a -> b) -> a -> b
$!
            if
              | w
w1 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW -> w
w2
              | w
w2 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
undefW -> w
w1
              | Bool
otherwise -> w -> w -> w
forall a. Ord a => a -> a -> a
min w
w1 w
w2

        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
wOld w -> w -> Bool
forall a. Eq a => a -> a -> Bool
/= w
undefW Bool -> Bool -> Bool
&& w
w' w -> w -> Bool
forall a. Ord a => a -> a -> Bool
< w
wOld) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState (ST s)) w
mat (Int -> Int -> Int
idx Int
from Int
to) w
w'
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
prev (Int -> Int -> Int
idx Int
from Int
to) (Int -> ST s ()) -> ST s Int -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
prev (Int -> Int -> Int
idx Int
b Int
to)
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
prev (Int -> Int -> Int
idx Int
from Int
b) Int
a
  where
    idx :: Int -> Int -> Int
idx !Int
from !Int
to = Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
from Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
to

-- | \(O(n)\) Given a predecessor array, retrieves a path from the root to a vertex.
--
-- ==== Constraints
-- - The path must not make a cycle, otherwise this function loops forever.
-- - There must be a path from the root to the @end@ vertex, otherwise the returned path is not
-- connected to the root.
--
-- @since 1.2.4.0
{-# INLINE constructPathFromRoot #-}
constructPathFromRoot :: (HasCallStack) => VU.Vector Int -> Int -> VU.Vector Int
constructPathFromRoot :: HasCallStack => Vector Int -> Int -> Vector Int
constructPathFromRoot Vector Int
parents = Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.reverse (Vector Int -> Vector Int)
-> (Int -> Vector Int) -> Int -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => Vector Int -> Int -> Vector Int
Vector Int -> Int -> Vector Int
constructPathToRoot Vector Int
parents

-- | \(O(n)\) Given a predecessor array, retrieves a path from a vertex to the root.
--
-- ==== Constraints
-- - The path must not make a cycle, otherwise this function loops forever.
-- - There must be a path from the root to the @end@ vertex, otherwise the returned path is not
-- connected to the root.
--
-- @since 1.2.4.0
{-# INLINEABLE constructPathToRoot #-}
constructPathToRoot :: (HasCallStack) => VU.Vector Int -> Int -> VU.Vector Int
constructPathToRoot :: HasCallStack => Vector Int -> Int -> Vector Int
constructPathToRoot Vector Int
parents = (Int -> Maybe (Int, Int)) -> Int -> Vector Int
forall a b. Unbox a => (b -> Maybe (a, b)) -> b -> Vector a
VU.unfoldr Int -> Maybe (Int, Int)
f
  where
    f :: Int -> Maybe (Int, Int)
f (-1) = Maybe (Int, Int)
forall a. Maybe a
Nothing
    f Int
v = (Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
v, Vector Int
parents Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v)

-- | \(O(n)\) Given a NxN predecessor matrix (created with `trackingFloydWarshall`), retrieves a
-- path from the root to an end vertex.
--
-- ==== Constraints
-- - The path must not make a cycle, otherwise this function loops forever.
-- - There must be a path from the root to the @end@ vertex, otherwise the returned path is not
-- connected to the root.
--
-- @since 1.2.4.0
{-# INLINE constructPathFromRootMat #-}
constructPathFromRootMat ::
  (HasCallStack) =>
  -- | Predecessor matrix.
  VU.Vector Int ->
  -- | Start vertex.
  Int ->
  -- | End vertex.
  Int ->
  -- | Path.
  VU.Vector Int
constructPathFromRootMat :: HasCallStack => Vector Int -> Int -> Int -> Vector Int
constructPathFromRootMat Vector Int
parents Int
start = Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.reverse (Vector Int -> Vector Int)
-> (Int -> Vector Int) -> Int -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => Vector Int -> Int -> Int -> Vector Int
Vector Int -> Int -> Int -> Vector Int
constructPathToRootMat Vector Int
parents Int
start

-- | \(O(n)\) Given a NxN predecessor matrix(created with `trackingFloydWarshall`), retrieves a
-- path from a vertex to the root.
--
-- ==== Constraints
-- - The path must not make a cycle, otherwise this function loops forever.
-- - There must be a path from the root to the @end@ vertex, otherwise the returned path is not
-- connected to the root.
--
-- @since 1.2.4.0
{-# INLINEABLE constructPathToRootMat #-}
constructPathToRootMat ::
  (HasCallStack) =>
  -- | Predecessor matrix.
  VU.Vector Int ->
  -- | Start vertex.
  Int ->
  -- | End vertex.
  Int ->
  -- | Path.
  VU.Vector Int
constructPathToRootMat :: HasCallStack => Vector Int -> Int -> Int -> Vector Int
constructPathToRootMat Vector Int
parents Int
start Int
end =
  let parents' :: Vector Int
parents' = Int -> Vector Int -> Vector Int
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take Int
n (Vector Int -> Vector Int) -> Vector Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ Int -> Vector Int -> Vector Int
forall a. Unbox a => Int -> Vector a -> Vector a
VU.drop (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
start) Vector Int
parents
   in HasCallStack => Vector Int -> Int -> Vector Int
Vector Int -> Int -> Vector Int
constructPathToRoot Vector Int
parents' Int
end
  where
    -- Assuming `n < 2^32`, it should always be correct:
    -- https://zenn.dev/mod_poppo/articles/atcoder-beginner-contest-284-d#%E8%A7%A3%E6%B3%953%EF%BC%9Asqrt%E3%81%A8round%E3%82%92%E4%BD%BF%E3%81%86
    Int
n :: Int = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double -> Int) -> (Double -> Double) -> Double -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
parents) :: Double)

-- | \(O(n)\) Given a NxN predecessor matrix (created with `newTrackingFloydWarshall`), retrieves a
-- path from the root to an end vertex.
--
-- ==== Constraints
-- - The path must not make a cycle, otherwise this function loops forever.
-- - There must be a path from the root to the @end@ vertex, otherwise the returned path is not
-- connected to the root.
--
-- @since 1.2.4.0
{-# INLINE constructPathFromRootMatM #-}
constructPathFromRootMatM ::
  (HasCallStack, PrimMonad m) =>
  -- | Predecessor matrix.
  VUM.MVector (PrimState m) Int ->
  -- | Start vertex.
  Int ->
  -- | End vertex.
  Int ->
  -- | Path.
  m (VU.Vector Int)
constructPathFromRootMatM :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MVector (PrimState m) Int -> Int -> Int -> m (Vector Int)
constructPathFromRootMatM MVector (PrimState m) Int
parents Int
start = (Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.reverse <$>) (m (Vector Int) -> m (Vector Int))
-> (Int -> m (Vector Int)) -> Int -> m (Vector Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MVector (PrimState m) Int -> Int -> Int -> m (Vector Int)
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MVector (PrimState m) Int -> Int -> Int -> m (Vector Int)
constructPathToRootMatM MVector (PrimState m) Int
parents Int
start

-- | \(O(n)\) Given a NxN predecessor matrix (created with `newTrackingFloydWarshall`), retrieves a
-- path from a vertex to the root.
--
-- ==== Constraints
-- - The path must not make a cycle, otherwise this function loops forever.
-- - There must be a path from the root to the @end@ vertex, otherwise the returned path is not
-- connected to the root.
--
-- @since 1.2.4.0
{-# INLINEABLE constructPathToRootMatM #-}
constructPathToRootMatM ::
  (HasCallStack, PrimMonad m) =>
  -- | Predecessor matrix.
  VUM.MVector (PrimState m) Int ->
  -- | Start vertex.
  Int ->
  -- | End vertex.
  Int ->
  -- | Path.
  m (VU.Vector Int)
constructPathToRootMatM :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MVector (PrimState m) Int -> Int -> Int -> m (Vector Int)
constructPathToRootMatM MVector (PrimState m) Int
parents Int
start Int
end = ST (PrimState m) (Vector Int) -> m (Vector Int)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Vector Int) -> m (Vector Int))
-> ST (PrimState m) (Vector Int) -> m (Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  Vector Int
parents' <- MVector (PrimState (ST (PrimState m))) Int
-> ST (PrimState m) (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) Int
MVector (PrimState (ST (PrimState m))) Int
parents
  Vector Int -> ST (PrimState m) (Vector Int)
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector Int -> ST (PrimState m) (Vector Int))
-> Vector Int -> ST (PrimState m) (Vector Int)
forall a b. (a -> b) -> a -> b
$ HasCallStack => Vector Int -> Int -> Int -> Vector Int
Vector Int -> Int -> Int -> Vector Int
constructPathToRootMat Vector Int
parents' Int
start Int
end