{-# 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@ 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,
    findCycleDirected,
    findCycleUndirected,

    -- * Generic graph functions

    -- TODO: generalize vertex dimensions?
    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

    -- TODO: panic instead of infinite loop?

    -- *** Single source 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\).
    constructPathFromRootMat,
    constructPathToRootMat,
    constructPathFromRootMatM,
    constructPathToRootMatM,
  )
where

import AtCoder.Dsu qualified as Dsu
import AtCoder.Extra.HashMap qualified as HM
import AtCoder.Extra.IntSet qualified as IS
import AtCoder.Extra.Ix0 (Bounds0, Ix0 (..))
import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.Internal.Buffer qualified as B
import AtCoder.Internal.Csr as Csr
import AtCoder.Internal.GrowVec qualified as GV
import AtCoder.Internal.MinHeap qualified as MH
import AtCoder.Internal.Queue qualified as Q
import AtCoder.Internal.Scc qualified as ACISCC
import Control.Applicative ((<|>))
import Control.Monad (replicateM_, unless, 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.Bits ((.<<.), (.|.))
import Data.Foldable (for_)
import Data.Maybe (fromJust, fromMaybe)
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 Data.Word (Word8)
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
wCsr :: forall w. Csr w -> Vector w
adjCsr :: forall w. Csr w -> Vector Int
startCsr :: forall w. Csr w -> Vector Int
mCsr :: forall w. Csr w -> Int
nCsr :: forall w. Csr w -> Int
..} = 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

-- TODO: is this minimum cycle?

-- | \(O(n + m)\) Given a directed graph, finds a minimal cycle and returns a vector of vertices and
-- a vector of @(vertices, csrEdgeIndices)@.
--
-- ==== __Example__
--
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let gr = Gr.build' 4 $ VU.fromList [(0, 1), (1, 2), (2, 3), (3, 1)]
-- >>> findCycleDirected gr -- returns (vs, es)
-- Just ([1,2,3],[1,2,3])
--
-- @since 1.4.0.0
{-# INLINEABLE findCycleDirected #-}
findCycleDirected :: (HasCallStack, VU.Unbox w) => Csr w -> Maybe (VU.Vector Int, VU.Vector Int)
findCycleDirected :: forall w.
(HasCallStack, Unbox w) =>
Csr w -> Maybe (Vector Int, Vector Int)
findCycleDirected gr :: Csr w
gr@Csr {Int
Vector w
Vector Int
wCsr :: forall w. Csr w -> Vector w
adjCsr :: forall w. Csr w -> Vector Int
startCsr :: forall w. Csr w -> Vector Int
mCsr :: forall w. Csr w -> Int
nCsr :: forall w. Csr w -> Int
nCsr :: Int
mCsr :: Int
startCsr :: Vector Int
adjCsr :: Vector Int
wCsr :: Vector w
..} = (forall s. ST s (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe (Vector Int, Vector Int)))
 -> Maybe (Vector Int, Vector Int))
-> (forall s. ST s (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  MVector s Word8
used <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate @_ @Word8 Int
nCsr Word8
0
  -- par <- VUM.unsafeNew @_ @(Int, Int) nCsr
  MVector s (Int, Int)
par <- Int -> (Int, Int) -> ST s (MVector (PrimState (ST s)) (Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nCsr (-Int
1 :: Int, -Int
1 :: Int)
  GrowVec s Int
vs <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (GrowVec (PrimState m) a)
GV.new @_ @Int Int
16
  GrowVec s Int
es <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (GrowVec (PrimState m) a)
GV.new @_ @Int Int
16
  GrowVec s Int
esFrom <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (GrowVec (PrimState m) a)
GV.new @_ @Int Int
16 -- If we had `from` in Csr, we could skip this
  let dfs :: Int -> m ()
dfs Int
u = do
        MVector (PrimState m) Word8 -> Int -> Word8 -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Word8
MVector (PrimState m) Word8
used Int
u Word8
1
        let next :: Vector (Int, Int) -> m ()
next Vector (Int, Int)
evs = case Vector (Int, Int) -> Maybe ((Int, Int), Vector (Int, Int))
forall a. Unbox a => Vector a -> Maybe (a, Vector a)
VU.uncons Vector (Int, Int)
evs of
              Maybe ((Int, Int), Vector (Int, Int))
Nothing -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              Just ((!Int
iEdge, !Int
v), !Vector (Int, Int)
evs') -> do
                Bool
b <- GrowVec (PrimState m) Int -> m Bool
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m Bool
GV.null GrowVec s Int
GrowVec (PrimState m) Int
es
                Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
                  Word8
use <- MVector (PrimState m) Word8 -> Int -> m Word8
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Word8
MVector (PrimState m) Word8
used Int
v
                  case Word8
use of
                    Word8
0 -> do
                      MVector (PrimState m) (Int, Int) -> Int -> (Int, Int) -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s (Int, Int)
MVector (PrimState m) (Int, Int)
par Int
v (Int
u, Int
iEdge)
                      Int -> m ()
dfs Int
v
                      Vector (Int, Int) -> m ()
next Vector (Int, Int)
evs'
                    Word8
1 -> do
                      GrowVec (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState m) Int
es Int
iEdge
                      GrowVec (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState m) Int
esFrom Int
u
                      let backtrack :: Int -> f ()
backtrack Int
cur
                            | Int
cur Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
v = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                            | Bool
otherwise = do
                                (!Int
prevVert, !Int
edge) <- MVector (PrimState f) (Int, Int) -> Int -> f (Int, Int)
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s (Int, Int)
MVector (PrimState f) (Int, Int)
par Int
cur
                                GrowVec (PrimState f) Int -> Int -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState f) Int
es Int
edge
                                GrowVec (PrimState f) Int -> Int -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState f) Int
esFrom Int
prevVert
                                Int -> f ()
backtrack Int
prevVert
                      Int -> m ()
forall {f :: * -> *}. (PrimState f ~ s, PrimMonad f) => Int -> f ()
backtrack Int
u
                      GrowVec (PrimState m) Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Ord a, Unbox a) =>
GrowVec (PrimState m) a -> m ()
GV.reverse GrowVec s Int
GrowVec (PrimState m) Int
es
                      GrowVec (PrimState m) Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Ord a, Unbox a) =>
GrowVec (PrimState m) a -> m ()
GV.reverse GrowVec s Int
GrowVec (PrimState m) Int
esFrom
                    Word8
_ -> do
                      Vector (Int, Int) -> m ()
next Vector (Int, Int)
evs'

        Vector (Int, Int) -> m ()
next (Vector (Int, Int) -> m ()) -> Vector (Int, Int) -> m ()
forall a b. (a -> b) -> a -> b
$ Csr w -> Int -> Vector (Int, Int)
forall w. HasCallStack => Csr w -> Int -> Vector (Int, Int)
eAdj Csr w
gr Int
u
        MVector (PrimState m) Word8 -> Int -> Word8 -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Word8
MVector (PrimState m) Word8
used Int
u Word8
2

  MVector (PrimState (ST s)) Word8
-> (Int -> Word8 -> ST s ()) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a b.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (Int -> a -> m b) -> m ()
VGM.iforM_ MVector s Word8
MVector (PrimState (ST s)) Word8
used ((Int -> Word8 -> ST s ()) -> ST s ())
-> (Int -> Word8 -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v Word8
use -> do
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word8
use Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      Int -> ST s ()
forall {f :: * -> *}. (PrimState f ~ s, PrimMonad f) => Int -> f ()
dfs Int
v

  Bool
b <- GrowVec (PrimState (ST s)) Int -> ST s Bool
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m Bool
GV.null GrowVec s Int
GrowVec (PrimState (ST s)) Int
es
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
b (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    -- find minimum cycle
    MVector s Int
nxt <- 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
nCsr (-Int
1 :: Int) -- edge indices
    do
      Vector Int
es' <- GrowVec (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowVec s Int
GrowVec (PrimState (ST s)) Int
es
      Vector Int
esFrom' <- GrowVec (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowVec s Int
GrowVec (PrimState (ST s)) Int
esFrom
      Vector (Int, Int) -> ((Int, Int) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Vector Int -> Vector Int -> Vector (Int, Int)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Int
es' Vector Int
esFrom') (((Int, Int) -> ST s ()) -> ST s ())
-> ((Int, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
iEdge, !Int
vFrom) -> 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
nxt Int
vFrom Int
iEdge

    [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
nCsr 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
vA -> do
      Int
nxtA <- 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
nxt Int
vA
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
nxtA Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        Vector (Int, Int) -> ((Int, Int) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Csr w -> Int -> Vector (Int, Int)
forall w. HasCallStack => Csr w -> Int -> Vector (Int, Int)
eAdj Csr w
gr Int
vA) (((Int, Int) -> ST s ()) -> ST s ())
-> ((Int, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
iEdge, !Int
vB) -> do
          Int
nxtB <- 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
nxt Int
vB
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
nxtB Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1 Bool -> Bool -> Bool
|| Vector Int
adjCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
nxtA Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
vB) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            let inner :: Int -> f ()
inner Int
x
                  | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
vB = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                  | Bool
otherwise = do
                      Int
nxtX <- MVector (PrimState f) Int -> Int -> Int -> f Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector s Int
MVector (PrimState f) Int
nxt Int
x (-Int
1)
                      Int -> f ()
inner (Int -> f ()) -> Int -> f ()
forall a b. (a -> b) -> a -> b
$ Vector Int
adjCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
nxtX
            Int -> ST s ()
forall {f :: * -> *}. (PrimState f ~ s, PrimMonad f) => Int -> f ()
inner Int
vA
            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
nxt Int
vA Int
iEdge

    GrowVec (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m ()
GV.clear GrowVec s Int
GrowVec (PrimState (ST s)) Int
es
    let loop :: Int -> f ()
loop Int
v
          | Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nCsr = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          | Bool
otherwise = do
              Int
nxtV <- MVector (PrimState f) Int -> Int -> f 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 f) Int
nxt Int
v
              if Int
nxtV Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1
                then Int -> f ()
loop (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                else do
                  let inner :: Int -> m ()
inner Int
x = do
                        GrowVec (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState m) Int
vs Int
x
                        Int
nxtX <- MVector (PrimState m) Int -> Int -> m 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 m) Int
nxt Int
x
                        GrowVec (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState m) Int
es Int
nxtX
                        let !x' :: Int
x' = Vector Int
adjCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
nxtX
                        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
x' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
v) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> m ()
inner Int
x'
                  Int -> f ()
forall {f :: * -> *}. (PrimState f ~ s, PrimMonad f) => Int -> f ()
inner Int
v
    Int -> ST s ()
forall {f :: * -> *}. (PrimState f ~ s, PrimMonad f) => Int -> f ()
loop Int
0

  Vector Int
vs' <- GrowVec (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowVec s Int
GrowVec (PrimState (ST s)) Int
vs
  Vector Int
es' <- GrowVec (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowVec s Int
GrowVec (PrimState (ST s)) Int
es
  if Vector Int -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector Int
es'
    then Maybe (Vector Int, Vector Int)
-> ST s (Maybe (Vector Int, Vector Int))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Vector Int, Vector Int)
forall a. Maybe a
Nothing
    else Maybe (Vector Int, Vector Int)
-> ST s (Maybe (Vector Int, Vector Int))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Vector Int, Vector Int)
 -> ST s (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
-> ST s (Maybe (Vector Int, Vector Int))
forall a b. (a -> b) -> a -> b
$ (Vector Int, Vector Int) -> Maybe (Vector Int, Vector Int)
forall a. a -> Maybe a
Just (Vector Int
vs', Vector Int
es')

-- | \(O(n + m)\) Given an undirected graph, finds a minimal cycle and returns a vector of vertices
-- a vector of @(vertices, csrEdgeIndices)@. A single edge index does not make much sense for an
-- undirected graph, so map back to the original edge index manually if needed.
--
-- ==== Constraints
-- - The graph must be created with `swapDupe` or `swapDupe'`. Otherwise the returned edge indices
--   could make no sense.
--
-- ==== __Example__
--
-- >>> import AtCoder.Extra.Graph qualified as Gr
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let gr = Gr.build' 4 . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (1, 3), (2, 3)]
-- >>> findCycleUndirected gr -- returns (vs, es)
-- Just ([1,3,2],[3,5,2])
--
-- Retrieve original edge indices that makes up the cycle, by recording them in edge weights:
--
-- >>> let gr = Gr.build 4 . Gr.swapDupe $ VU.fromList [(0, 1, 0 :: Int), (1, 2, 1), (1, 3, 2), (2, 3, 3)]
-- >>> let Just (vs, es) = findCycleUndirected gr -- returns (vs, es)
-- >>> VU.backpermute (Gr.wCsr gr) es
-- [2,3,1]
--
-- It's a bit hacky.
--
-- @since 1.4.0.0
{-# INLINEABLE findCycleUndirected #-}
findCycleUndirected :: (HasCallStack, VU.Unbox w) => Csr w -> Maybe (VU.Vector Int, VU.Vector Int)
findCycleUndirected :: forall w.
(HasCallStack, Unbox w) =>
Csr w -> Maybe (Vector Int, Vector Int)
findCycleUndirected gr :: Csr w
gr@Csr {Int
Vector w
Vector Int
wCsr :: forall w. Csr w -> Vector w
adjCsr :: forall w. Csr w -> Vector Int
startCsr :: forall w. Csr w -> Vector Int
mCsr :: forall w. Csr w -> Int
nCsr :: forall w. Csr w -> Int
nCsr :: Int
mCsr :: Int
startCsr :: Vector Int
adjCsr :: Vector Int
wCsr :: Vector w
..} =
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int -> Bool
forall a. Integral a => a -> Bool
even Int
mCsr) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Graph.findCycleUndirected: the number of edge in an undirected graph must be even: `" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
mCsr String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"`"
   in -- If we have the same edge id for duplicated edges, `findCycleSimpleUndirected` could be modified
      -- to handle both complex and simple graph. We don't, and we need the complex graph handling.
      -- This is not optimal, but we need a dedicated `buildUndirected` function and different edge ID
      -- (not index) handling in CSR if we go with the optimal approach.
      --
      --  Note that the implementations are suspecious..
      Csr w -> Maybe (Vector Int, Vector Int)
forall w.
(HasCallStack, Unbox w) =>
Csr w -> Maybe (Vector Int, Vector Int)
findCycleComplexUndirected Csr w
gr Maybe (Vector Int, Vector Int)
-> Maybe (Vector Int, Vector Int) -> Maybe (Vector Int, Vector Int)
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Csr w -> Maybe (Vector Int, Vector Int)
forall w.
(HasCallStack, Unbox w) =>
Csr w -> Maybe (Vector Int, Vector Int)
findCycleSimpleUndirected Csr w
gr

{-# INLINEABLE findCycleComplexUndirected #-}
findCycleComplexUndirected :: (HasCallStack, VU.Unbox w) => Csr w -> Maybe (VU.Vector Int, VU.Vector Int)
findCycleComplexUndirected :: forall w.
(HasCallStack, Unbox w) =>
Csr w -> Maybe (Vector Int, Vector Int)
findCycleComplexUndirected gr :: Csr w
gr@Csr {Int
Vector w
Vector Int
wCsr :: forall w. Csr w -> Vector w
adjCsr :: forall w. Csr w -> Vector Int
startCsr :: forall w. Csr w -> Vector Int
mCsr :: forall w. Csr w -> Int
nCsr :: forall w. Csr w -> Int
nCsr :: Int
mCsr :: Int
startCsr :: Vector Int
adjCsr :: Vector Int
wCsr :: Vector w
..} = (forall s. ST s (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe (Vector Int, Vector Int)))
 -> Maybe (Vector Int, Vector Int))
-> (forall s. ST s (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  HashMap s Int
usedE <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (HashMap (PrimState m) a)
HM.new @_ @Int (Int
mCsr Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ {- not needed, but in case of panic? -} Int
4)
  HashMap s Word8
cntE <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (HashMap (PrimState m) a)
HM.new @_ @Word8 (Int
mCsr Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ {- not needed, but in case of panic? -} Int
4)

  -- we'll give unique indices to (u, v) pairs
  let ix :: a -> a -> a
ix a
u a
v = a -> a -> a
forall a. Ord a => a -> a -> a
min a
u a
v a -> Int -> a
forall a. Bits a => a -> Int -> a
.<<. Int
32 a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a -> a
forall a. Ord a => a -> a -> a
max a
u a
v

  let nextU :: Int -> f (Maybe (Vector Int, Vector Int))
nextU Int
u
        | Int
u Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nCsr = Maybe (Vector Int, Vector Int)
-> f (Maybe (Vector Int, Vector Int))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Vector Int, Vector Int)
forall a. Maybe a
Nothing
        | Bool
otherwise = do
            let nextV :: Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int))
nextV Vector (Int, Int)
evs = case Vector (Int, Int) -> Maybe ((Int, Int), Vector (Int, Int))
forall a. Unbox a => Vector a -> Maybe (a, Vector a)
VU.uncons Vector (Int, Int)
evs of
                  Maybe ((Int, Int), Vector (Int, Int))
Nothing -> Maybe (Vector Int, Vector Int)
-> f (Maybe (Vector Int, Vector Int))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Vector Int, Vector Int)
forall a. Maybe a
Nothing
                  Just ((!Int
e, !Int
v), !Vector (Int, Int)
evs') -> case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
u Int
v of
                    -- self loop edge
                    Ordering
EQ -> Maybe (Vector Int, Vector Int)
-> f (Maybe (Vector Int, Vector Int))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Vector Int, Vector Int)
 -> f (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
-> f (Maybe (Vector Int, Vector Int))
forall a b. (a -> b) -> a -> b
$ (Vector Int, Vector Int) -> Maybe (Vector Int, Vector Int)
forall a. a -> Maybe a
Just (Int -> Vector Int
forall a. Unbox a => a -> Vector a
VU.singleton Int
v, Int -> Vector Int
forall a. Unbox a => a -> Vector a
VU.singleton Int
e)
                    Ordering
LT -> do
                      let !i :: Int
i = Int -> Int -> Int
forall {a}. (Bits a, Ord a) => a -> a -> a
ix Int
u Int
v
                      Word8
c <- Word8 -> Maybe Word8 -> Word8
forall a. a -> Maybe a -> a
fromMaybe Word8
0 (Maybe Word8 -> Word8) -> f (Maybe Word8) -> f Word8
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap (PrimState f) Word8 -> Int -> f (Maybe Word8)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap s Word8
HashMap (PrimState f) Word8
cntE Int
i
                      case Word8
c of
                        Word8
0 -> do
                          HashMap (PrimState f) Int -> Int -> Int -> f ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap s Int
HashMap (PrimState f) Int
usedE Int
i Int
e
                          HashMap (PrimState f) Word8 -> Int -> Word8 -> f ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap s Word8
HashMap (PrimState f) Word8
cntE Int
i Word8
1
                          Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int))
nextV Vector (Int, Int)
evs'
                        Word8
1 -> do
                          -- found the first duplicated edge
                          HashMap (PrimState f) Word8 -> Int -> Word8 -> f ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap s Word8
HashMap (PrimState f) Word8
cntE Int
i Word8
2
                          Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int))
nextV Vector (Int, Int)
evs'
                        Word8
_ -> do
                          Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int))
nextV Vector (Int, Int)
evs'
                    Ordering
GT -> do
                      let !i :: Int
i = Int -> Int -> Int
forall {a}. (Bits a, Ord a) => a -> a -> a
ix Int
u Int
v
                      Word8
cnt <- Word8 -> Maybe Word8 -> Word8
forall a. a -> Maybe a -> a
fromMaybe Word8
0 (Maybe Word8 -> Word8) -> f (Maybe Word8) -> f Word8
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap (PrimState f) Word8 -> Int -> f (Maybe Word8)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap s Word8
HashMap (PrimState f) Word8
cntE Int
i
                      case Word8
cnt of
                        Word8
2 -> do
                          -- there are duplicate edges between (u, v) and this is the
                          -- first (u, v) (u > v)
                          HashMap (PrimState f) Word8 -> Int -> Word8 -> f ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap s Word8
HashMap (PrimState f) Word8
cntE Int
i Word8
3
                          Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int))
nextV Vector (Int, Int)
evs'
                        Word8
3 -> do
                          -- this is the second duplicate edge (u, v) (u > v)
                          Int
e1 <- Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> f (Maybe Int) -> f Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap (PrimState f) Int -> Int -> f (Maybe Int)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap s Int
HashMap (PrimState f) Int
usedE Int
i
                          let vs :: Vector Int
vs = Int -> [Int] -> Vector Int
forall a. Unbox a => Int -> [a] -> Vector a
VU.fromListN Int
2 [Int
v, Int
u]
                          let es :: Vector Int
es = Int -> [Int] -> Vector Int
forall a. Unbox a => Int -> [a] -> Vector a
VU.fromListN Int
2 [Int
e1, Int
e]
                          Maybe (Vector Int, Vector Int)
-> f (Maybe (Vector Int, Vector Int))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Vector Int, Vector Int)
 -> f (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
-> f (Maybe (Vector Int, Vector Int))
forall a b. (a -> b) -> a -> b
$ (Vector Int, Vector Int) -> Maybe (Vector Int, Vector Int)
forall a. a -> Maybe a
Just (Vector Int
vs, Vector Int
es)
                        Word8
_ -> Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int))
nextV Vector (Int, Int)
evs'

            Maybe (Vector Int, Vector Int)
res <- Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int))
forall {f :: * -> *}.
(PrimState f ~ s, PrimMonad f) =>
Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int))
nextV (Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int)))
-> Vector (Int, Int) -> f (Maybe (Vector Int, Vector Int))
forall a b. (a -> b) -> a -> b
$ Csr w -> Int -> Vector (Int, Int)
forall w. HasCallStack => Csr w -> Int -> Vector (Int, Int)
eAdj Csr w
gr Int
u
            case Maybe (Vector Int, Vector Int)
res of
              Just (Vector Int, Vector Int)
ret -> Maybe (Vector Int, Vector Int)
-> f (Maybe (Vector Int, Vector Int))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Vector Int, Vector Int)
 -> f (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
-> f (Maybe (Vector Int, Vector Int))
forall a b. (a -> b) -> a -> b
$ (Vector Int, Vector Int) -> Maybe (Vector Int, Vector Int)
forall a. a -> Maybe a
Just (Vector Int, Vector Int)
ret
              Maybe (Vector Int, Vector Int)
Nothing -> Int -> f (Maybe (Vector Int, Vector Int))
nextU (Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

  Int -> ST s (Maybe (Vector Int, Vector Int))
forall {f :: * -> *}.
(PrimState f ~ s, PrimMonad f) =>
Int -> f (Maybe (Vector Int, Vector Int))
nextU Int
0

{-# INLINEABLE findCycleSimpleUndirected #-}
findCycleSimpleUndirected :: (HasCallStack, VU.Unbox w) => Csr w -> Maybe (VU.Vector Int, VU.Vector Int)
findCycleSimpleUndirected :: forall w.
(HasCallStack, Unbox w) =>
Csr w -> Maybe (Vector Int, Vector Int)
findCycleSimpleUndirected gr :: Csr w
gr@Csr {Int
Vector w
Vector Int
wCsr :: forall w. Csr w -> Vector w
adjCsr :: forall w. Csr w -> Vector Int
startCsr :: forall w. Csr w -> Vector Int
mCsr :: forall w. Csr w -> Int
nCsr :: forall w. Csr w -> Int
nCsr :: Int
mCsr :: Int
startCsr :: Vector Int
adjCsr :: Vector Int
wCsr :: Vector w
..} = (forall s. ST s (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe (Vector Int, Vector Int)))
 -> Maybe (Vector Int, Vector Int))
-> (forall s. ST s (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  -- marks both (u, v) and (v, u)
  HashMap s Bit
usedUV <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (HashMap (PrimState m) a)
HM.new @_ @Bit (Int
mCsr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4)

  -- we'll give unique indices to (u, v) pairs
  let ix :: a -> a -> a
ix a
u a
v = a -> a -> a
forall a. Ord a => a -> a -> a
min a
u a
v a -> Int -> a
forall a. Bits a => a -> Int -> a
.<<. Int
32 a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a -> a
forall a. Ord a => a -> a -> a
max a
u a
v

  -- depth
  MVector s Int
dep <- 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
nCsr (-Int
1 :: Int)

  -- vertex -> edge index
  MVector s Int
par <- 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
nCsr (-Int
1 :: Int)
  MVector s Int
parFrom <- 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
nCsr (-Int
1 :: Int)

  -- Get DFS forest
  let dfs :: Int -> Int -> m ()
dfs Int
u Int
d = do
        MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState m) Int
dep Int
u Int
d
        Vector (Int, Int) -> ((Int, Int) -> m ()) -> m ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Csr w -> Int -> Vector (Int, Int)
forall w. HasCallStack => Csr w -> Int -> Vector (Int, Int)
eAdj Csr w
gr Int
u) (((Int, Int) -> m ()) -> m ()) -> ((Int, Int) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(!Int
iEdge, !Int
v) -> do
          Int
dv <- MVector (PrimState m) Int -> Int -> m 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 m) Int
dep Int
v
          Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
dv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
            -- we're marking both direction of an undirected edge
            HashMap (PrimState m) Bit -> Int -> Bit -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap s Bit
HashMap (PrimState m) Bit
usedUV (Int -> Int -> Int
forall {a}. (Bits a, Ord a) => a -> a -> a
ix Int
u Int
v) (Bit -> m ()) -> Bit -> m ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
            MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState m) Int
par Int
v Int
iEdge
            MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState m) Int
parFrom Int
v Int
u
            Int -> Int -> m ()
dfs Int
v (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

  MVector (PrimState (ST s)) Int
-> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a b.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (Int -> a -> m b) -> m ()
VGM.iforM_ MVector s Int
MVector (PrimState (ST s)) Int
dep ((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
1) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      Int -> Int -> ST s ()
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
Int -> Int -> m ()
dfs Int
v Int
0

  GrowVec s Int
vs <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (GrowVec (PrimState m) a)
GV.new @_ @Int Int
16
  GrowVec s Int
es <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (GrowVec (PrimState m) a)
GV.new @_ @Int Int
16
  Vector Int
dep' <- 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
dep

  -- Find edge with minimum depth difference, which makes up a loop (not used in the DFS forets):
  MVector s Int
minLen <- 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
forall a. Bounded a => a
maxBound Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 :: Int)
  MVector s (Int, Int)
backE <- Int -> (Int, Int) -> ST s (MVector (PrimState (ST s)) (Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 (-Int
1 :: Int, -Int
1 :: 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
nCsr 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
vA -> do
    let !dA :: Int
dA = Vector Int
dep' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
vA
    Vector (Int, Int) -> ((Int, Int) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Csr w -> Int -> Vector (Int, Int)
forall w. HasCallStack => Csr w -> Int -> Vector (Int, Int)
eAdj Csr w
gr Int
vA) (((Int, Int) -> ST s ()) -> ST s ())
-> ((Int, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
iEdge, !Int
vB) -> do
      Bool
b <- Bool -> (Bit -> Bool) -> Maybe Bit -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False Bit -> Bool
unBit (Maybe Bit -> Bool) -> ST s (Maybe Bit) -> ST s Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap (PrimState (ST s)) Bit -> Int -> ST s (Maybe Bit)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap s Bit
HashMap (PrimState (ST s)) Bit
usedUV (Int -> Int -> Int
forall {a}. (Bits a, Ord a) => a -> a -> a
ix Int
vA Int
vB)
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
b (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        let !dB :: Int
dB = Vector Int
dep' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
vB
        let !d :: Int
d = Int -> Int
forall a. Num a => a -> a
abs (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int
dA Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
dB
        Int
minLen' <- 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
minLen Int
0
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
minLen') (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
minLen Int
0 Int
d
          MVector (PrimState (ST s)) (Int, Int)
-> 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, Int)
MVector (PrimState (ST s)) (Int, Int)
backE Int
0 (Int
iEdge, Int
vA)

  (!Int
backE', !Int
backFrom) <- MVector (PrimState (ST s)) (Int, Int) -> Int -> ST s (Int, Int)
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s (Int, Int)
MVector (PrimState (ST s)) (Int, Int)
backE Int
0
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
backE' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= -Int
1) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    let try :: Int -> Int -> m ()
try Int
a Int
b = do
          if Vector Int
dep' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Vector Int
dep' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
b
            then Int -> Int -> m ()
try Int
b Int
a
            else do
              -- v_1 -> v_N -> v_{N - 1} -> .. -> v_2 -> v_1
              GrowVec (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState m) Int
es Int
backE'
              GrowVec (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState m) Int
vs Int
a
              let backtrack :: Int -> f ()
backtrack Int
v = do
                    Bool -> f () -> f ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
a) (f () -> f ()) -> f () -> f ()
forall a b. (a -> b) -> a -> b
$ do
                      Int
parE <- MVector (PrimState f) Int -> Int -> f 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 f) Int
par Int
v
                      Int
v' <- MVector (PrimState f) Int -> Int -> f 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 f) Int
parFrom Int
v
                      GrowVec (PrimState f) Int -> Int -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState f) Int
vs Int
v
                      GrowVec (PrimState f) Int -> Int -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
GV.pushBack GrowVec s Int
GrowVec (PrimState f) Int
es Int
parE
                      Int -> f ()
backtrack Int
v'
              Int -> m ()
forall {f :: * -> *}. (PrimState f ~ s, PrimMonad f) => Int -> f ()
backtrack Int
b
    Int -> Int -> ST s ()
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
Int -> Int -> m ()
try Int
backFrom (Vector Int
adjCsr Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
backE')

  Vector Int
vs' <- GrowVec (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowVec s Int
GrowVec (PrimState (ST s)) Int
vs
  Vector Int
es' <- GrowVec (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowVec s Int
GrowVec (PrimState (ST s)) Int
es
  if Vector Int -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector Int
es'
    then Maybe (Vector Int, Vector Int)
-> ST s (Maybe (Vector Int, Vector Int))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Vector Int, Vector Int)
forall a. Maybe a
Nothing
    else Maybe (Vector Int, Vector Int)
-> ST s (Maybe (Vector Int, Vector Int))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Vector Int, Vector Int)
 -> ST s (Maybe (Vector Int, Vector Int)))
-> Maybe (Vector Int, Vector Int)
-> ST s (Maybe (Vector Int, Vector Int))
forall a b. (a -> b) -> a -> b
$ (Vector Int, Vector Int) -> Maybe (Vector Int, Vector Int)
forall a. a -> Maybe a
Just (Vector Int
vs', Vector Int
es')

-- -------------------------------------------------------------------------------------------------
-- 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; there must be no cycle.
--
-- ==== __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 ::
  -- | \(n\): The number of vertices.
  Int ->
  -- | \(g\): Graph function, typically @'adj' gr@.
  (Int -> VU.Vector Int) ->
  -- | Vertices in topological ordering: upstream vertices come first.
  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: both \((u, v)\) and \((v, u)\) edges must exist.
--
-- ==== __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 ::
  -- | \(n\): The number of vertices.
  Int ->
  -- | \(g\): Graph function, typically @'adj' gr@.
  (Int -> VU.Vector Int) ->
  -- | Connected components.
  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 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 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 -> m b
dfs !b
acc Int
u = do
        Bit Bool
b <- MVector (PrimState m) Bit -> Int -> Bit -> m Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector s Bit
MVector (PrimState m) Bit
vis Int
u (Bit -> m Bit) -> Bit -> m Bit
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
        if Bool
b
          then b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
acc
          else do
            Buffer (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer s Int
Buffer (PrimState m) Int
buf Int
u
            (b -> Int -> m b) -> b -> Vector Int -> m b
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM' b -> Int -> m 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 {m :: * -> *} {b}.
(PrimState m ~ s, PrimMonad m, Num b) =>
b -> Int -> m 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 s Int
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 ::
  -- | \(n\): The number of vertices.
  Int ->
  -- | \(g\): Graph function, typically @'adj' gr@.
  (Int -> VU.Vector Int) ->
  -- | Bipartite vertex coloring.
  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 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 s
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 s
Dsu (PrimState (ST s))
dsu Int
u (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)

      MVector 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 s
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 s Bit
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 s Bit
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 s Bit
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 s
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 s Bit
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 s Bit
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 s
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 s Bit -> MVector s Bit
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VGM.take Int
n MVector s Bit
color
      let isCompatible :: Int -> f Bool
isCompatible Int
v
            | Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Bool -> f Bool
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
            | Bool
otherwise = do
                Bit
c1 <- MVector (PrimState f) Bit -> Int -> f Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Bit
MVector (PrimState f) Bit
color (Int -> f Bit) -> f Int -> f Bit
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Dsu (PrimState f) -> Int -> f Int
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
Dsu (PrimState m) -> Int -> m Int
Dsu.leader Dsu s
Dsu (PrimState f)
dsu Int
v
                Bit
c2 <- MVector (PrimState f) Bit -> Int -> f Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Bit
MVector (PrimState f) Bit
color (Int -> f Bit) -> f Int -> f Bit
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Dsu (PrimState f) -> Int -> f Int
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
Dsu (PrimState m) -> Int -> m Int
Dsu.leader Dsu s
Dsu (PrimState f)
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 -> f Bool
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
                  else Int -> f Bool
isCompatible (Int -> f Bool) -> Int -> f 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
forall {f :: * -> *}.
(PrimState f ~ s, PrimMonad f) =>
Int -> f 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)

-- | \(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 ::
  -- | \(n\): The number of vertices.
  Int ->
  -- | \(g\): Graph function, typically @'adj' gr@.
  (Int -> VU.Vector Int) ->
  -- | Graph that represents a block-cut tree, where super vertices \((n \ge n)\) represent each
  -- biconnected component.
  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 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 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 s Int
st <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Buffer (PrimState m) a)
B.new @_ @Int Int
n
  MVector 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 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 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 -> m Int
dfs Int
k0 Int
v Int
p = do
        Buffer (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer s Int
Buffer (PrimState m) Int
st Int
v
        MVector (PrimState m) Bit -> Int -> Bit -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Bit
MVector (PrimState m) Bit
used Int
v (Bit -> m ()) -> Bit -> m ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
        MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState m) Int
low Int
v Int
k0
        MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState m) Int
ord Int
v Int
k0

        (Int, Int) -> Int
forall a b. (a, b) -> b
snd
          ((Int, Int) -> Int) -> m (Int, Int) -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Int, Int) -> Int -> m (Int, Int))
-> (Int, Int) -> Vector Int -> m (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) -> m (Int, Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
child, Int
k)
                  else do
                    Bit Bool
b <- MVector (PrimState m) Bit -> Int -> m Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Bit
MVector (PrimState m) 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 m) Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m Int
B.length Buffer s Int
Buffer (PrimState m) Int
st
                        Int
k' <- Int -> Int -> Int -> m Int
dfs Int
k Int
to Int
v
                        Int
lowTo <- MVector (PrimState m) Int -> Int -> m 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 m) Int
low Int
to
                        MVector (PrimState m) Int -> (Int -> Int) -> Int -> m ()
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 m) Int
low (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
lowTo) Int
v
                        Int
ordV <- MVector (PrimState m) Int -> Int -> m 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 m) Int
ord Int
v
                        Bool -> m () -> m ()
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)) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
                          Int
nxt <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState m) Int
next Int
0
                          MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s Int
MVector (PrimState m) Int
next Int
0 (Int
nxt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                          Buffer (PrimState m) (Int, Int) -> (Int, Int) -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer s (Int, Int)
Buffer (PrimState m) (Int, Int)
edges (Int
nxt, Int
v)
                          Int
len <- Buffer (PrimState m) Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m Int
B.length Buffer s Int
Buffer (PrimState m) Int
st
                          Int -> m () -> m ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
                            Int
back <- Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> m (Maybe Int) -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Buffer (PrimState m) Int -> m (Maybe Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Maybe a)
B.popBack Buffer s Int
Buffer (PrimState m) Int
st
                            Buffer (PrimState m) (Int, Int) -> (Int, Int) -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer s (Int, Int)
Buffer (PrimState m) (Int, Int)
edges (Int
nxt, Int
back)
                        (Int, Int) -> m (Int, Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
child', Int
k')
                      else do
                        Int
ordTo <- MVector (PrimState m) Int -> Int -> m 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 m) Int
ord Int
to
                        MVector (PrimState m) Int -> (Int -> Int) -> Int -> m ()
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 m) Int
low (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
ordTo) Int
v
                        (Int, Int) -> m (Int, Int)
forall a. a -> m 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
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
Int -> Int -> Int -> m 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 s Int
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 s Int
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 s Int
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 s (Int, Int)
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 s Int
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 s Bit
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 s Int
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 s (Int, Int)
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 ::
  -- | \(n\): The number of vertices.
  Int ->
  -- | \(g\): Graph function, typically @'adj' gr@.
  (Int -> VU.Vector Int) ->
  -- | Block-cut components
  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 function 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 function 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 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 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 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 s Int
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 s Int
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 s (i, Int)
Queue (PrimState (ST s)) (i, Int)
deque (i
src, Int
w0)

      let step :: i -> Int -> m ()
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 m) Int -> Int -> m 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 m) Int
dist Int
i0
            Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
w0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
wReserved0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
              Vector (i, Int) -> ((i, Int) -> m ()) -> m ()
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) -> m ()) -> m ()) -> ((i, Int) -> m ()) -> m ()
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 m) Int -> Int -> m 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 m) Int
dist Int
i
                -- NOTE: Do pruning just like Dijkstra:
                Bool -> m () -> m ()
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) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
                  MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState m) Int
dist Int
i Int
w
                  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
                    MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState m) Int
prev Int
i Int
i0
                  if Int
dw Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                    then Queue (PrimState m) (i, Int) -> (i, Int) -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
Q.pushFront Queue s (i, Int)
Queue (PrimState m) (i, Int)
deque (i
vExt, Int
w)
                    else Queue (PrimState m) (i, Int) -> (i, Int) -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
Q.pushBack Queue s (i, Int)
Queue (PrimState m) (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 s (i, Int)
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 ()
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
i -> Int -> m ()
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 s Int
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 s Int
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 initial 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 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)

  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 s w
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 s w
MVector (PrimState (ST s)) w
dist Int
v w
w
  MVector 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 -> m ()
update Int
v1 = do
        w
d1 <- MVector (PrimState m) w -> Int -> m 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 m) w
dist Int
v1
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (w
d1 w -> w -> Bool
forall a. Eq a => a -> a -> Bool
/= w
undefW) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
          Vector (Int, w) -> ((Int, w) -> m ()) -> m ()
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) -> m ()) -> m ()) -> ((Int, w) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(!Int
v2, !w
dw) -> do
            w
d2 <- MVector (PrimState m) w -> Int -> m 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 m) w
dist Int
v2
            let !d2' :: w
d2' = w
d1 w -> w -> w
forall a. Num a => a -> a -> a
+ w
dw
            Bool -> m () -> m ()
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) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
              MVector (PrimState m) w -> Int -> w -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s w
MVector (PrimState m) w
dist Int
v2 w
d2'
              Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
trackPrev (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
                MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState m) 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 m) Bool -> Int -> Bool -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Bool
MVector (PrimState m) Bool
updated Int
0 Bool
True

  let runLoop :: Int -> f (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) -> f (Maybe (Vector w, Vector Int))
forall a. a -> f 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 -> f ()) -> f ()
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 -> f ()
forall {m :: * -> *}. (PrimState m ~ s, PrimMonad m) => Int -> m ()
update
            Bool
b <- MVector (PrimState f) Bool -> Int -> Bool -> f Bool
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector s Bool
MVector (PrimState f) Bool
updated Int
0 Bool
False
            if Bool
b
              then Int -> f (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))
-> f (Vector w, Vector Int) -> f (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))
-> f (Vector w) -> f (Vector Int -> (Vector w, Vector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState f) w -> f (Vector w)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s w
MVector (PrimState f) w
dist f (Vector Int -> (Vector w, Vector Int))
-> f (Vector Int) -> f (Vector w, Vector Int)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState f) Int -> f (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState f) Int
prev)

  Int -> ST s (Maybe (Vector w, Vector Int))
forall {f :: * -> *}.
(PrimState f ~ s, PrimMonad f) =>
Int -> f (Maybe (Vector w, Vector Int))
runLoop Int
0

-- | \(O(n^3)\) Floyd–Warshall algorithm that returns a distance matrix \(m\).
--
-- - The distance matrix should be accessed as @m VG.! (`index0` (n, n) (from, to))@,
-- - There's a negative loop if there's any vertex \(v\) where @m VU.! (`index0` (n, n) (v, v))@
-- is negative.
--
-- ==== __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 VG.! (`index0` (n, n) (from, to))@,
-- - The predecessor matrix should be accessed as @m VG.! (`index0` (n, n) (root, v))@
-- - There's a negative loop if there's any vertex \(v\) where @m VU.! (`index0` (n, n) (v, v))@
-- is negative.
--
-- ==== __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\).
--
-- - The distance matrix should be accessed as @m VG.! (`index0` (n, n) (from, to))@,
-- - There's a negative loop if there's any vertex \(v\) where @m VU.! (`index0` (n, n) (v, v))@
-- 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.
--
-- - The distance matrix should be accessed as @m VG.! (`index0` (n, n) (from, to))@,
-- - The predecessor matrix should be accessed as @m VG.! (`index0` (n, n) (root, v))@
-- - There's a negative loop if there's any vertex \(v\) where @m VU.! (`index0` (n, n) (v, v))@
-- 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 change 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 chaneg 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 @sink@ 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 @sink@ 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 a sink vertex.
--
-- ==== Constraints
-- - The path must not make a cycle, otherwise this function loops forever.
-- - There must be a path from the root to the @sink@ 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 ->
  -- | Source vertex.
  Int ->
  -- | Sink vertex.
  Int ->
  -- | Path.
  VU.Vector Int
constructPathFromRootMat :: HasCallStack => Vector Int -> Int -> Int -> Vector Int
constructPathFromRootMat Vector Int
parents Int
source = 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
source

-- | \(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 @sink@ 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 ->
  -- | Source vertex.
  Int ->
  -- | Sink vertex.
  Int ->
  -- | Path.
  VU.Vector Int
constructPathToRootMat :: HasCallStack => Vector Int -> Int -> Int -> Vector Int
constructPathToRootMat Vector Int
parents Int
source Int
sink =
  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
source) Vector Int
parents
   in HasCallStack => Vector Int -> Int -> Vector Int
Vector Int -> Int -> Vector Int
constructPathToRoot Vector Int
parents' Int
sink
  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 a sink vertex.
--
-- ==== Constraints
-- - The path must not make a cycle, otherwise this function loops forever.
-- - There must be a path from the root to the @nd@ 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 ->
  -- | Source vertex.
  Int ->
  -- | Sink 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
source = (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
source

-- | \(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 @sink@ 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 ->
  -- | Source vertex.
  Int ->
  -- | Sink 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
source Int
sink = 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
source Int
sink