module AtCoder.Extra.Graph
(
module Csr,
swapDupe,
swapDupe',
scc,
topSort,
)
where
import AtCoder.Extra.IntSet qualified as IS
import AtCoder.Internal.Buffer qualified as B
import AtCoder.Internal.Csr as Csr
import AtCoder.Internal.Scc qualified as ACISCC
import Control.Monad (when)
import Control.Monad.ST (runST)
import Data.Foldable (for_)
import Data.Vector qualified as V
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
{-# INLINE swapDupe #-}
swapDupe :: (VU.Unbox (Int, Int, w)) => VU.Vector (Int, Int, w) -> VU.Vector (Int, Int, w)
swapDupe :: forall w.
Unbox (Int, Int, w) =>
Vector (Int, Int, w) -> Vector (Int, Int, w)
swapDupe = ((Int, Int, w) -> Vector (Int, Int, w))
-> Vector (Int, Int, w) -> Vector (Int, Int, w)
forall a b.
(Unbox a, Unbox b) =>
(a -> Vector b) -> Vector a -> Vector b
VU.concatMap (\(!Int
u, !Int
v, !w
w) -> Int -> [(Int, Int, w)] -> Vector (Int, Int, w)
forall a. Unbox a => Int -> [a] -> Vector a
VU.fromListN Int
2 [(Int
u, Int
v, w
w), (Int
v, Int
u, w
w)])
{-# INLINE swapDupe' #-}
swapDupe' :: (VU.Unbox (Int, Int)) => VU.Vector (Int, Int) -> VU.Vector (Int, Int)
swapDupe' :: Unbox (Int, Int) => Vector (Int, Int) -> Vector (Int, Int)
swapDupe' = ((Int, Int) -> Vector (Int, Int))
-> Vector (Int, Int) -> Vector (Int, Int)
forall a b.
(Unbox a, Unbox b) =>
(a -> Vector b) -> Vector a -> Vector b
VU.concatMap (\(!Int
u, !Int
v) -> Int -> [(Int, Int)] -> Vector (Int, Int)
forall a. Unbox a => Int -> [a] -> Vector a
VU.fromListN Int
2 [(Int
u, Int
v), (Int
v, Int
u)])
{-# 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
{-# INLINE topSort #-}
topSort :: Int -> (Int -> VU.Vector Int) -> VU.Vector Int
topSort :: Int -> (Int -> Vector Int) -> Vector Int
topSort Int
n Int -> Vector Int
gr = (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int)) -> Vector Int)
-> (forall s. ST s (Vector Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
MVector s Int
inDeg <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
n (Int
0 :: Int)
[Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
u -> do
Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ (Int -> Vector Int
gr Int
u) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v -> do
MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
inDeg (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
v
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
let run :: ST s ()
run = 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 ()
run
ST s ()
run
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