ac-library-hs 1.1.0.0 → 1.1.1.0
raw patch · 26 files changed
+1641/−292 lines, 26 filesPVP: major bump suggested
API removals or changes: PVP suggests a major version bump
API changes (from Hackage documentation)
+ AtCoder.Extra.Graph: blockCut :: Int -> (Int -> Vector Int) -> Csr ()
+ AtCoder.Extra.Graph: blockCutComponents :: Int -> (Int -> Vector Int) -> Vector (Vector Int)
+ AtCoder.Extra.Semigroup.Matrix: detMint :: forall a. KnownNat a => Matrix (ModInt a) -> ModInt a
+ AtCoder.Extra.Semigroup.Matrix: detMod :: Int -> Matrix Int -> Int
+ AtCoder.Extra.Semigroup.Matrix: inv :: forall a. (Fractional a, Eq a, Unbox a) => Matrix a -> Maybe (a, Matrix a)
+ AtCoder.Extra.Semigroup.Matrix: invRaw :: forall a. (Fractional a, Eq a, Unbox a) => Matrix a -> Maybe (a, Vector (Vector a))
+ AtCoder.Extra.Semigroup.Matrix: rank :: (Fractional a, Eq a, Unbox a) => Matrix a -> Int
+ AtCoder.Extra.Semigroup.Matrix: square :: (HasCallStack, Unbox a) => Int -> Vector a -> Matrix a
+ AtCoder.Extra.Tree.Lct: Lct :: {-# UNPACK #-} !Int -> !MVector s Vertex -> !MVector s Vertex -> !MVector s Vertex -> !MVector s Int -> !MVector s Bit -> !MVector s a -> !MVector s a -> !MVector s a -> !MVector s a -> !MVector s a -> !a -> a -> Lct s a
+ AtCoder.Extra.Tree.Lct: [dualProdLct] :: Lct s a -> !MVector s a
+ AtCoder.Extra.Tree.Lct: [invOpLct] :: Lct s a -> !a -> a
+ AtCoder.Extra.Tree.Lct: [lLct] :: Lct s a -> !MVector s Vertex
+ AtCoder.Extra.Tree.Lct: [midLct] :: Lct s a -> !MVector s a
+ AtCoder.Extra.Tree.Lct: [nLct] :: Lct s a -> {-# UNPACK #-} !Int
+ AtCoder.Extra.Tree.Lct: [pLct] :: Lct s a -> !MVector s Vertex
+ AtCoder.Extra.Tree.Lct: [prodLct] :: Lct s a -> !MVector s a
+ AtCoder.Extra.Tree.Lct: [rLct] :: Lct s a -> !MVector s Vertex
+ AtCoder.Extra.Tree.Lct: [revLct] :: Lct s a -> !MVector s Bit
+ AtCoder.Extra.Tree.Lct: [sLct] :: Lct s a -> !MVector s Int
+ AtCoder.Extra.Tree.Lct: [subtreeProdLct] :: Lct s a -> !MVector s a
+ AtCoder.Extra.Tree.Lct: [vLct] :: Lct s a -> !MVector s a
+ AtCoder.Extra.Tree.Lct: build :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Vector a -> Vector (Vertex, Vertex) -> m (Lct (PrimState m) a)
+ AtCoder.Extra.Tree.Lct: buildInv :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => (a -> a) -> Vector a -> Vector (Vertex, Vertex) -> m (Lct (PrimState m) a)
+ AtCoder.Extra.Tree.Lct: cut :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m ()
+ AtCoder.Extra.Tree.Lct: data Lct s a
+ AtCoder.Extra.Tree.Lct: evert :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Vertex -> m ()
+ AtCoder.Extra.Tree.Lct: expose :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Vertex -> m Vertex
+ AtCoder.Extra.Tree.Lct: expose_ :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Vertex -> m ()
+ AtCoder.Extra.Tree.Lct: instance GHC.Classes.Eq AtCoder.Extra.Tree.Lct.NodePlaceLct
+ AtCoder.Extra.Tree.Lct: jump :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> Int -> m Vertex
+ AtCoder.Extra.Tree.Lct: lca :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Int -> Int -> m Vertex
+ AtCoder.Extra.Tree.Lct: link :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m ()
+ AtCoder.Extra.Tree.Lct: modify :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> (a -> a) -> Vertex -> m ()
+ AtCoder.Extra.Tree.Lct: modifyM :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> (a -> m a) -> Vertex -> m ()
+ AtCoder.Extra.Tree.Lct: new :: (PrimMonad m, Monoid a, Unbox a) => Int -> m (Lct (PrimState m) a)
+ AtCoder.Extra.Tree.Lct: newInv :: (PrimMonad m, Monoid a, Unbox a) => (a -> a) -> Int -> m (Lct (PrimState m) a)
+ AtCoder.Extra.Tree.Lct: parent :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Int -> m (Maybe Vertex)
+ AtCoder.Extra.Tree.Lct: prodPath :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m a
+ AtCoder.Extra.Tree.Lct: prodSubtree :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m a
+ AtCoder.Extra.Tree.Lct: root :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Int -> m Vertex
+ AtCoder.Extra.Tree.Lct: type Vertex = Int
+ AtCoder.Extra.Tree.Lct: write :: (HasCallStack, PrimMonad m, Monoid a, Unbox a) => Lct (PrimState m) a -> Vertex -> a -> m ()
+ AtCoder.Internal.Buffer: popBack_ :: (PrimMonad m, Unbox a) => Buffer (PrimState m) a -> m ()
- AtCoder.Extra.Semigroup.Matrix: mul :: (Num e, Unbox e) => Matrix e -> Matrix e -> Matrix e
+ AtCoder.Extra.Semigroup.Matrix: mul :: forall e. (Num e, Unbox e) => Matrix e -> Matrix e -> Matrix e
Files
- CHANGELOG.md +7/−0
- README.md +1/−1
- ac-library-hs.cabal +9/−7
- benchmarks/Bench/Matrix.hs +15/−5
- benchmarks/BenchLib/Matrix.hs +134/−93
- benchmarks/Main.hs +1/−1
- src/AtCoder/Convolution.hs +3/−0
- src/AtCoder/Dsu.hs +1/−1
- src/AtCoder/Extra/Graph.hs +117/−1
- src/AtCoder/Extra/IntervalMap.hs +59/−34
- src/AtCoder/Extra/MultiSet.hs +12/−0
- src/AtCoder/Extra/Pdsu.hs +35/−14
- src/AtCoder/Extra/Semigroup/Matrix.hs +277/−45
- src/AtCoder/Extra/Tree/Hld.hs +10/−10
- src/AtCoder/Extra/Tree/Lct.hs +844/−0
- src/AtCoder/Extra/Tree/TreeMonoid.hs +15/−5
- src/AtCoder/Extra/WaveletMatrix.hs +18/−18
- src/AtCoder/Extra/WaveletMatrix/Raw.hs +26/−26
- src/AtCoder/Internal/Buffer.hs +10/−0
- src/AtCoder/Internal/Convolution.hs +6/−6
- src/AtCoder/Internal/Math.hs +5/−1
- src/AtCoder/Internal/Scc.hs +2/−2
- src/AtCoder/Internal/String.hs +3/−3
- src/AtCoder/String.hs +3/−3
- test/Tests/Extra/HashMap.hs +14/−15
- test/Tests/Extra/Semigroup/Matrix.hs +14/−1
CHANGELOG.md view
@@ -1,5 +1,12 @@ # Revision history for acl-hs +## 1.1.1.0 -- Jan 2025++- Added `AtCoder.Extra.Tree.Lct`+- Added `blockCut`, `blockCutComponents` in `AtCoder.Extra.Graph`+- Added `popBack_` in `AtCoder.Internal.Buffer`+- Added `square`, `rank`, `inv`, `invRaw`, `detMod`, `detMint` in `AtCoder.Extra.Matrix`+ ## 1.1.0.0 -- Jan 2025 - Removed `RangeSetId` and `RangeAddId` from `AtCoder.Extra.Monoid`.
README.md view
@@ -5,7 +5,7 @@ ## Notes - The library is mainly for AtCoder and only GHC 9.8.4 is guaranteed to be supported.-- Functions primarily use half-open interval \([l, r)\).+- Functions primarily use half-open interval [l, r). - The extra module contains additional utilities beyond the original C++ library. ## Usage
ac-library-hs.cabal view
@@ -4,7 +4,7 @@ -- PVP summary: +-+------- breaking API changes -- | | +----- non-breaking API additions -- | | | +--- code changes with no API change-version: 1.1.0.0+version: 1.1.1.0 synopsis: Data structures and algorithms description: Haskell port of [ac-library](https://github.com/atcoder/ac-library), a library for competitive@@ -74,14 +74,11 @@ AtCoder.Convolution AtCoder.Dsu AtCoder.Extra.Bisect+ AtCoder.Extra.Graph AtCoder.Extra.HashMap AtCoder.Extra.IntervalMap AtCoder.Extra.IntMap AtCoder.Extra.IntSet- AtCoder.Extra.Graph- AtCoder.Extra.Tree- AtCoder.Extra.Tree.Hld- AtCoder.Extra.Tree.TreeMonoid AtCoder.Extra.Math AtCoder.Extra.Monoid AtCoder.Extra.Monoid.Affine1@@ -91,9 +88,13 @@ AtCoder.Extra.Monoid.RollingHash AtCoder.Extra.Monoid.V2 AtCoder.Extra.MultiSet+ AtCoder.Extra.Pdsu AtCoder.Extra.Semigroup.Matrix AtCoder.Extra.Semigroup.Permutation- AtCoder.Extra.Pdsu+ AtCoder.Extra.Tree+ AtCoder.Extra.Tree.Hld+ AtCoder.Extra.Tree.Lct+ AtCoder.Extra.Tree.TreeMonoid AtCoder.Extra.WaveletMatrix AtCoder.Extra.WaveletMatrix.BitVector AtCoder.Extra.WaveletMatrix.Raw@@ -245,10 +246,11 @@ , hspec , mtl , QuickCheck- -- , quickcheck-instances , random , tagged , tasty , tasty-hunit , tasty-quickcheck , tasty-rerun++-- , quickcheck-instances
benchmarks/Bench/Matrix.hs view
@@ -1,8 +1,8 @@ module Bench.Matrix (benches) where -import AtCoder.ModInt qualified as M-import AtCoder.Extra.Semigroup.Matrix qualified as ACMAT import AtCoder.Extra.Math qualified as ACEM+import AtCoder.Extra.Semigroup.Matrix qualified as ACMAT+import AtCoder.ModInt qualified as M import BenchLib.Matrix qualified as Mat import Control.Monad.State.Class (MonadState, state) import Control.Monad.Trans.State.Strict (evalState, runState)@@ -43,17 +43,24 @@ bench "mul3_1" $ whnf (V.foldl1' Mat.mul3_1) randomMatrixInput, bench "mul3_2" $ whnf (V.foldl1' Mat.mul3_2) randomMatrixInput, bench "mul3_3" $ whnf (V.foldl1' Mat.mul3_3) randomMatrixInput,+ bench "mul4_1" $ whnf (V.foldl1' Mat.mul4_1) randomMatrixInput,+ bench "mul4_2" $ whnf (V.foldl1' Mat.mul4_2) randomMatrixInput,+ bench "mul4_3" $ whnf (V.foldl1' Mat.mul4_3) randomMatrixInput, bench "mulMod1" $ whnf (V.foldl1' (Mat.mulMod1 m)) randomMatrixInput, bench "mulMod2" $ whnf (V.foldl1' (Mat.mulMod2 m)) randomMatrixInput, bench "mulMod3" $ whnf (V.foldl1' (Mat.mulMod3 m)) randomMatrixInput, bench "mulMod4" $ whnf (V.foldl1' (Mat.mulMod4 m)) randomMatrixInput, bench "mulMod5" $ whnf (V.foldl1' (Mat.mulMod5 m)) randomMatrixInput,+ bench "mulMod6" $ whnf (V.foldl1' (Mat.mulMod6 m)) randomMatrixInput,+ bench "mulMod7" $ whnf (V.foldl1' (Mat.mulMod7 m)) randomMatrixInput, bench "mulMint1" $ whnf (V.foldl1' Mat.mulMint1) randomMintMatrixInput, bench "mulMint2" $ whnf (V.foldl1' Mat.mulMint2) randomMintMatrixInput, bench "mulMint3" $ whnf (V.foldl1' Mat.mulMint3) randomMintMatrixInput,+ bench "mulMint4" $ whnf (V.foldl1' Mat.mulMint4) randomMintMatrixInput, -- mul (ACL) bench "mul_ACL" $ whnf (V.foldl1' ACMAT.mul) randomMatrixInputACL, bench "mulMod_ACL" $ whnf (V.foldl1' (ACMAT.mulMod m)) randomMatrixInputACL,+ bench "mulMint_ACL" $ whnf (V.foldl1' ACMAT.mulMint) randomMintMatrixInputACL, -- pow mod (ACL only) bench "powMod_ACL" $ whnf (VU.foldl' (flip (ACMAT.powMod m)) squareMat) randomVec, bench "powMintACL_stimes" $ whnf (VU.foldl' (flip stimes) squareMatMint) randomVec,@@ -72,16 +79,19 @@ -- ACL matrix randomMatrixInputACL :: V.Vector (ACMAT.Matrix Int)- !randomMatrixInputACL = V.map (\mat -> ACMAT.new (Mat.hM mat) (Mat.wM mat) (Mat.vecM mat)) randomMatrixInput+ !randomMatrixInputACL = V.map (\(Mat.Matrix h w vec) -> ACMAT.new h w vec) randomMatrixInput + randomMintMatrixInputACL :: V.Vector (ACMAT.Matrix (M.ModInt 998244353))+ !randomMintMatrixInputACL = V.map (\(Mat.Matrix h w vec) -> ACMAT.new h w vec) randomMintMatrixInput+ squareMat :: ACMAT.Matrix Int !squareMat = evalState (randomSquareACLMatrix 17) $ mkStdGen 123456789 - squareMatMint ::ACMAT.Matrix (M.ModInt 998244353)+ squareMatMint :: ACMAT.Matrix (M.ModInt 998244353) !squareMatMint = ACMAT.map M.new squareMat -- non-zero random vector randomVec :: VU.Vector Int !randomVec = VU.map ((+ 1) . fromIntegral) $- VU.unfoldrExactN 100 (genWord64R (m - 2)) (mkStdGen 123456789)+ VU.unfoldrExactN 100 (genWord64R (998244353 - 2)) (mkStdGen 123456789)
benchmarks/BenchLib/Matrix.hs view
@@ -12,23 +12,32 @@ mul3_1, mul3_2, mul3_3,+ mul4_1,+ mul4_2,+ mul4_3, mulMod1, mulMod2, mulMod3, mulMod4, mulMod5,+ mulMod6,+ mulMod7, mulMint1, mulMint2, mulMint3,+ mulMint4, ) where import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Barrett qualified as BT import AtCoder.ModInt qualified as M+import Data.Foldable (for_) 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 (Word64) import GHC.Exts (Proxy#, proxy#) import GHC.Stack (HasCallStack)@@ -73,7 +82,7 @@ {-# INLINE mul1 #-} mul1 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e-mul1 !a !b =+mul1 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -86,19 +95,13 @@ (0, 0) where f row col = VU.sum $ VU.zipWith (*) (rows1 VG.! row) (cols2 VG.! col)- h = hM a- w = wM a- vecA = vecM a- h' = hM b- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" rows1 = V.unfoldrExactN h (VU.splitAt w) vecA cols2 = V.generate w' $ \col -> VU.unfoldrExactN h' (\i -> (VG.unsafeIndex vecB i, i + w')) col {-# INLINE mul2 #-} mul2 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e-mul2 !a !b =+mul2 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -111,18 +114,12 @@ (0, 0) where f row col = VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (rows1 VG.! row)- h = hM a- w = wM a- vecA = vecM a- h' = hM b- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" rows1 = V.unfoldrExactN h (VU.splitAt w) vecA {-# INLINE mul3_1 #-} mul3_1 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e-mul3_1 !a !b =+mul3_1 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -135,17 +132,11 @@ (0, 0) where f row col = VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (VU.unsafeSlice (w * row) w vecA)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" {-# INLINE mul3_2 #-} mul3_2 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e-mul3_2 !a !b =+mul3_2 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -157,31 +148,64 @@ ) (0, 0) where- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" {-# INLINE mul3_3 #-} mul3_3 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e-mul3_3 !a !b = Matrix h w' $ VU.generate (h * w') $ \i ->+mul3_3 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.generate (h * w') $ \i -> let (!row, !col) = i `quotRem` w' in VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (VU.unsafeSlice (w * row) w vecA) where- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" +-- | Fastest: efficient iteration+{-# INLINE mul4_1 #-}+mul4_1 :: forall e. (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul4_1 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.create $ do+ c <- VUM.replicate (h * w') (0 :: e)+ for_ [0 .. h - 1] $ \i -> do+ for_ [0 .. w - 1] $ \k -> do+ for_ [0 .. w' - 1] $ \j -> do+ let !aik = VG.unsafeIndex vecA (i * w + k)+ let !bkj = VG.unsafeIndex vecB (k * w' + j)+ VGM.unsafeModify c (+ (aik * bkj)) (i * w' + j)+ pure c+ where+ !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mul4_2 #-}+mul4_2 :: forall e. (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul4_2 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.create $ do+ c <- VUM.replicate (h * w') (0 :: e)+ for_ [0 .. h - 1] $ \i -> do+ let !iw = i * w+ let !iw' = i * w'+ for_ [0 .. w - 1] $ \k -> do+ let !kw' = k * w'+ for_ [0 .. w' - 1] $ \j -> do+ let !aik = VG.unsafeIndex vecA (iw + k)+ let !bkj = VG.unsafeIndex vecB (kw' + j)+ VGM.unsafeModify c (+ (aik * bkj)) (iw' + j)+ pure c+ where+ !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mul4_3 #-}+mul4_3 :: forall e. (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul4_3 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.create $ do+ c <- VUM.replicate (h * w') (0 :: e)+ for_ [0 .. h - 1] $ \i -> do+ for_ [0 .. w - 1] $ \k -> do+ VU.iforM_ (VU.unsafeSlice (k * w') w' vecB) $ \j bkj -> do+ let !aik = VG.unsafeIndex vecA (i * w + k)+ VGM.unsafeModify c (+ (aik * bkj)) (i * w' + j)+ pure c+ where+ !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"+ {-# INLINE mulMod1 #-} mulMod1 :: Int -> Matrix Int -> Matrix Int -> Matrix Int-mulMod1 !m !a !b =+mulMod1 !m (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -196,17 +220,11 @@ f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA) addMod x y = (x + y) `mod` m mulMod x y = (x * y) `mod` m- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" {-# INLINE mulMod2 #-} mulMod2 :: Int -> Matrix Int -> Matrix Int -> Matrix Int-mulMod2 !m !a !b =+mulMod2 !m (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -224,17 +242,11 @@ | x + y >= m = x + y - m | otherwise = x + y mulMod x y = (x * y) `mod` m- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" {-# INLINE mulMod3 #-} mulMod3 :: Int -> Matrix Int -> Matrix Int -> Matrix Int-mulMod3 !m !a !b =+mulMod3 !m (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -250,17 +262,11 @@ f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA) addMod x y = (x + y) `mod` m mulMod x y = fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" {-# INLINE mulMod4 #-} mulMod4 :: Int -> Matrix Int -> Matrix Int -> Matrix Int-mulMod4 !m !a !b =+mulMod4 !m (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -276,17 +282,11 @@ f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA) addMod x y = (x + y) `rem` m mulMod x y = fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" {-# INLINE mulMod5 #-} mulMod5 :: Int -> Matrix Int -> Matrix Int -> Matrix Int-mulMod5 !m !a !b =+mulMod5 !m (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -299,22 +299,58 @@ (0, 0) where !bt = BT.new32 $ fromIntegral m+ -- NOTE: this is unsafe if the matrix is too large+ f row col =+ fromIntegral+ . (`rem` fromIntegral m)+ . VU.sum+ $ VU.imap+ (\iRow x -> BT.mulMod bt (fromIntegral x) (fromIntegral (VG.unsafeIndex vecB (col + (iRow * w')))))+ (VU.unsafeSlice (w * row) w vecA)+ !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mulMod6 #-}+mulMod6 :: Int -> Matrix Int -> Matrix Int -> Matrix Int+mulMod6 !m (Matrix h w vecA) (Matrix h' w' vecB) =+ Matrix h w' $+ VU.unfoldrExactN+ (h * w')+ ( \(!row, !col) ->+ let !x = f row col+ in if col + 1 >= w'+ then (x, (row + 1, 0))+ else (x, (row, col + 1))+ )+ (0, 0)+ where+ !bt = BT.new32 $ fromIntegral m f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA) addMod x y | x + y >= m = x + y - m | otherwise = x + y mulMod x y = fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" +{-# INLINE mulMod7 #-}+mulMod7 :: Int -> Matrix Int -> Matrix Int -> Matrix Int+mulMod7 !m (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.create $ do+ c <- VUM.replicate (h * w') (0 :: Int)+ for_ [0 .. h - 1] $ \i -> do+ for_ [0 .. w - 1] $ \k -> do+ for_ [0 .. w' - 1] $ \j -> do+ let !aik = VG.unsafeIndex vecA (i * w + k)+ let !bkj = VG.unsafeIndex vecB (k * w' + j)+ VGM.unsafeModify c (addMod (mulMod_ aik bkj)) (i * w' + j)+ pure c+ where+ !bt = BT.new32 $ fromIntegral m+ addMod x y = (x + y) `rem` m+ mulMod_ x y = fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)+ !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMod: matrix size mismatch"+ {-# INLINE mulMint1 #-} mulMint1 :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)-mulMint1 !a !b =+mulMint1 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -330,17 +366,11 @@ f row col = VU.sum $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA) mulMod :: M.ModInt a -> M.ModInt a -> M.ModInt a mulMod = (*)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" {-# INLINE mulMint2 #-} mulMint2 :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)-mulMint2 !a !b =+mulMint2 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -357,18 +387,11 @@ f row col = VU.sum $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA) mulMod :: M.ModInt a -> M.ModInt a -> M.ModInt a mulMod (M.ModInt x) (M.ModInt y) = M.unsafeNew . fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" --- REMARK: This is very unsafe in that it can overflow (mod^2 * n) {-# INLINE mulMint3 #-} mulMint3 :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)-mulMint3 !a !b =+mulMint3 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -381,14 +404,32 @@ (0, 0) where !bt = BT.new32 $ fromIntegral (natVal' (proxy# @a))+ -- NOTE: this is unsafe if the matrix is too large f :: Int -> Int -> M.ModInt a- f row col = M.new64 . VU.sum $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+ f row col =+ M.new64+ . VU.sum+ $ VU.imap+ (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w'))))+ (VU.unsafeSlice (w * row) w vecA) mulMod :: M.ModInt a -> M.ModInt a -> Word64 mulMod (M.ModInt x) (M.ModInt y) = BT.mulMod bt (fromIntegral x) (fromIntegral y)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++-- performs memory efficient iteration, but requires more type conversions and slow+{-# INLINE mulMint4 #-}+mulMint4 :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)+mulMint4 (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.create $ do+ c <- VUM.replicate (h * w') (M.unsafeNew 0)+ for_ [0 .. h - 1] $ \i -> do+ for_ [0 .. w - 1] $ \k -> do+ for_ [0 .. w' - 1] $ \j -> do+ let !aik = VG.unsafeIndex vecA (i * w + k)+ let !bkj = VG.unsafeIndex vecB (k * w' + j)+ VGM.unsafeModify c (+ (mulMod_ aik bkj)) (i * w' + j)+ pure c+ where+ !bt = BT.new32 $ fromIntegral (natVal' (proxy# @a))+ mulMod_ :: M.ModInt a -> M.ModInt a -> M.ModInt a+ mulMod_ (M.ModInt x) (M.ModInt y) = M.unsafeNew . fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)+ !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMint: matrix size mismatch"
benchmarks/Main.hs view
@@ -4,8 +4,8 @@ import Bench.AddMod qualified import Bench.Matrix qualified-import Bench.MulMod qualified import Bench.ModInt qualified+import Bench.MulMod qualified import Bench.PowMod qualified import Criterion.Main
src/AtCoder/Convolution.hs view
@@ -76,6 +76,7 @@ -- - \(O(n\log{n} + \log{\mathrm{mod}})\), where \(n = |a| + |b|\). -- -- @since 1.0.0.0+{-# INLINABLE convolution #-} convolution :: forall p. (HasCallStack, AM.Modulus p) =>@@ -106,6 +107,7 @@ -- - \(O(n\log{n} + \log{\mathrm{mod}})\), where \(n = |a| + |b|\). -- -- @since 1.0.0.0+{-# INLINABLE convolutionRaw #-} convolutionRaw :: forall p a. (HasCallStack, AM.Modulus p, Integral a, VU.Unbox a) =>@@ -135,6 +137,7 @@ -- - \(O(n\log{n})\), where \(n = |a| + |b|\). -- -- @since 1.0.0.0+{-# INLINABLE convolution64 #-} convolution64 :: (HasCallStack) => VU.Vector Int ->
src/AtCoder/Dsu.hs view
@@ -46,7 +46,7 @@ ( -- * Disjoint set union Dsu (nDsu), - -- * Constructor+ -- * Constructors new, -- * Merging
src/AtCoder/Extra/Graph.hs view
@@ -14,6 +14,8 @@ -- * Graph search topSort,+ blockCut,+ blockCutComponents, ) where @@ -23,7 +25,9 @@ import AtCoder.Internal.Scc qualified as ACISCC import Control.Monad (when) import Control.Monad.ST (runST)+import Data.Bit (Bit (..)) import Data.Foldable (for_)+import Data.Maybe (fromJust) import Data.Vector qualified as V import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU@@ -100,6 +104,8 @@ scc :: Csr w -> V.Vector (VU.Vector Int) scc = ACISCC.sccCsr +-- TODO: change scc to take arbitrary graph form+ -- | \(O(n \log n + m)\) Returns the lexicographically smallest topological ordering of the given -- graph. --@@ -115,7 +121,7 @@ -- [1,2,4,0,3] -- -- @since 1.1.0.0-{-# INLINE topSort #-}+{-# INLINABLE topSort #-} topSort :: Int -> (Int -> VU.Vector Int) -> VU.Vector Int topSort n gr = runST $ do inDeg <- VUM.replicate n (0 :: Int)@@ -145,3 +151,113 @@ run B.unsafeFreeze buf++-- | \(O(n + m)\) Returns a [block cut tree](https://en.wikipedia.org/wiki/Biconnected_component)+-- where super vertices 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+{-# INLINABLE blockCut #-}+blockCut :: Int -> (Int -> VU.Vector Int) -> Csr ()+blockCut n gr = runST $ do+ low <- VUM.replicate n (0 :: Int)+ ord <- VUM.replicate n (0 :: Int)+ st <- B.new @_ @Int n+ used <- VUM.replicate n $ Bit False+ edges <- B.new @_ @(Int, Int {- TODO: correct capacity? -}) (2 * n)+ -- represents the bidirected component's index. also works as super vertex indices.+ next <- VUM.replicate 1 n++ let dfs k0 v p = do+ B.pushBack st v+ VGM.write used v $ Bit True+ VGM.write low v k0+ VGM.write ord v k0++ snd+ <$> VU.foldM'+ ( \(!child, !k) to -> do+ if to == p+ then pure (child, k)+ else do+ Bit b <- VGM.read used to+ if not b+ then do+ let !child' = child + 1+ s <- B.length st+ k' <- dfs k to v+ lowTo <- VGM.read low to+ VGM.modify low (min lowTo) v+ ordV <- VGM.read ord v+ when ((p == -1 && child' > 1) || (p /= -1 && lowTo >= ordV)) $ do+ nxt <- VGM.unsafeRead next 0+ VGM.unsafeWrite next 0 (nxt + 1)+ B.pushBack edges (nxt, v)+ len <- B.length st+ for_ [1 .. len - s] $ \_ -> do+ back <- fromJust <$> B.popBack st+ B.pushBack edges (nxt, back)+ pure (child', k')+ else do+ ordTo <- VGM.read ord to+ VGM.modify low (min ordTo) v+ pure (child, k)+ )+ (0 :: Int, k0 + 1)+ (gr v)++ _ <-+ VGM.ifoldM'+ ( \k v (Bit b) -> do+ if b+ then do+ pure k+ else do+ k' <- dfs k v (-1)+ st' <- B.unsafeFreeze st+ nxt <- VGM.unsafeRead next 0+ VGM.unsafeWrite next 0 (nxt + 1)+ VU.forM_ st' $ \x -> do+ B.pushBack edges (nxt, x)+ B.clear st+ pure k'+ )+ (0 :: Int)+ used++ n' <- VGM.unsafeRead next 0+ Csr.build' n' <$> B.unsafeFreeze edges++-- | \(O(n + m)\) Returns a [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+{-# INLINE blockCutComponents #-}+blockCutComponents :: Int -> (Int -> VU.Vector Int) -> V.Vector (VU.Vector Int)+blockCutComponents n gr =+ let bct = blockCut n gr+ d = nCsr bct - n+ in V.generate d ((bct `adj`) . (+ n))
src/AtCoder/Extra/IntervalMap.hs view
@@ -90,7 +90,7 @@ import AtCoder.Extra.IntMap qualified as IM import Control.Monad (foldM_)-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) import Data.Vector.Generic qualified as G import Data.Vector.Unboxed qualified as VU import GHC.Stack (HasCallStack)@@ -109,6 +109,7 @@ -- | \(O(n)\) Creates an empty `IntervalMap`. -- -- @since 1.1.0.0+{-# INLINE new #-} new :: (PrimMonad m, VU.Unbox a) => Int -> m (IntervalMap (PrimState m) a) new = fmap IntervalMap . IM.new @@ -121,6 +122,7 @@ -- [(0,(2,10)),(2,(4,11)),(4,(6,12))] -- -- @since 1.1.0.0+{-# INLINE build #-} build :: (PrimMonad m, Eq a, VU.Unbox a) => VU.Vector a -> m (IntervalMap (PrimState m) a) build xs = buildM xs onAdd where@@ -130,6 +132,7 @@ -- interval, while performing @onAdd@ hook for each interval. -- -- @since 1.1.0.0+{-# INLINE buildM #-} buildM :: (PrimMonad m, Eq a, VU.Unbox a) => -- | Input values@@ -159,6 +162,7 @@ -- | \(O(\log n)\) Returns whether a point \(x\) is contained within any of the intervals. -- -- @since 1.1.0.0+{-# INLINE contains #-} contains :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> m Bool contains itm i = intersects itm i (i + 1) @@ -166,6 +170,7 @@ -- intervals. -- -- @since 1.1.0.0+{-# INLINE intersects #-} intersects :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m Bool intersects (IntervalMap dim) l r | l >= r = pure False@@ -178,6 +183,7 @@ -- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\). -- -- @since 1.1.0.0+{-# INLINE lookup #-} lookup :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m (Maybe (Int, Int, a)) lookup (IntervalMap im) l r | l >= r = pure Nothing@@ -192,6 +198,7 @@ -- Throws an error if no such interval exists. -- -- @since 1.1.0.0+{-# INLINE read #-} read :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m a read itm l r = do res <- readMaybe itm l r@@ -203,6 +210,7 @@ -- Returns `Nothing` if no such interval exists. -- -- @since 1.1.0.0+{-# INLINE readMaybe #-} readMaybe :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m (Maybe a) readMaybe (IntervalMap dim) l r | l >= r = pure Nothing@@ -217,6 +225,7 @@ -- map. Overwrites any overlapping intervals. -- -- @since 1.1.0.0+{-# INLINE insert #-} insert :: (PrimMonad m, Eq a, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> a -> m () insert itm l r x = insertM itm l r x onAdd onDel where@@ -228,6 +237,7 @@ -- hooks. -- -- @since 1.1.0.0+{-# INLINABLE insertM #-} insertM :: (PrimMonad m, Eq a, VU.Unbox a) => -- | The map@@ -249,10 +259,10 @@ !r <- handleRight l0 r0 (!l', !r') <- handleLeft l0 r onAdd l' r' x- IM.insert dim l' (r', x)+ stToPrim $ IM.insert dim l' (r', x) where handleRight l r = do- res <- IM.lookupGE dim l+ res <- stToPrim $ IM.lookupGE dim l case res of Just interval0@(!_, (!_, !_)) -> run interval0 l r Nothing -> pure r@@ -270,7 +280,7 @@ | l' == r && x' == x = do -- adjacent interval with the same value: merge into one. onDel l' r' x'- IM.delete_ dim l'+ stToPrim $ IM.delete_ dim l' pure r' | l' == r = do -- adjacent interval with different values: nothing to do.@@ -279,8 +289,9 @@ | r' <= r = do -- inside the interval: delete and continue onDel l' r' x'- IM.delete_ dim l'- res <- IM.lookupGT dim l'+ res <- stToPrim $ do+ IM.delete_ dim l'+ IM.lookupGT dim l' case res of Just rng -> run rng l r Nothing -> pure r@@ -288,18 +299,20 @@ | x' == x = do -- intersecting interval with the same value: merge into one. onDel l' r' x'- IM.delete_ dim l'- pure r'+ stToPrim $ do+ IM.delete_ dim l'+ pure r' | otherwise = do -- intersecting interval with a different value: delete the intersection. onDel l' r' x' onAdd r r' x'- IM.delete_ dim l'- IM.insert dim r (r', x')- pure r+ stToPrim $ do+ IM.delete_ dim l'+ IM.insert dim r (r', x')+ pure r handleLeft l r = do- res <- IM.lookupLT dim l+ res <- stToPrim $ IM.lookupLT dim l case res of Nothing -> pure (l, r) Just (!l', (!r', !x'))@@ -307,8 +320,9 @@ | r' == l && x' == x -> do -- adjacent interval with the same value: merge into one. onDel l' r' x'- IM.delete_ dim l'- pure (l', r)+ stToPrim $ do+ IM.delete_ dim l'+ pure (l', r) | r' == l -> do -- adjacent interval with different values: nothing to do. pure (l, r)@@ -319,28 +333,32 @@ | x' == x -> do -- insersecting interval with the same value: merge into one. onDel l' r' x'- IM.delete_ dim l'- pure (min l l', max r r')+ stToPrim $ do+ IM.delete_ dim l'+ pure (min l l', max r r') | r' > r -> do -- [l', r') contains [l, r) with a different value: split into three. onDel l' r' x' onAdd l' l x' onAdd r r' x'- -- IM.delete_ dim l'- IM.insert dim l' (l, x')- IM.insert dim r (r', x')- pure (l, r)+ stToPrim $ do+ -- IM.delete_ dim l'+ IM.insert dim l' (l, x')+ IM.insert dim r (r', x')+ pure (l, r) | otherwise -> do -- insersecting interval with a different value: delete. onDel l' r' x' onAdd l' l x'- -- IM.delete_ dim l'- IM.insert dim l' (l, x')- pure (l, r)+ stToPrim $ do+ -- IM.delete_ dim l'+ IM.insert dim l' (l, x')+ pure (l, r) -- | Amortized \(O(\log n)\) Deletes an interval \([l, r)\) from the map. -- -- @since 1.1.0.0+{-# INLINE delete #-} delete :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m () delete itm l r = deleteM itm l r onAdd onDel where@@ -351,6 +369,7 @@ -- changes via @onAdd@ and @onDel@ hooks. -- -- @since 1.1.0.0+{-# INLINABLE deleteM #-} deleteM :: (PrimMonad m, VU.Unbox a) => -- | The map@@ -371,7 +390,7 @@ handleLeft l0 r0 where handleRight l r = do- res <- IM.lookupGE dim l+ res <- stToPrim $ IM.lookupGE dim l case res of Just interval0@(!_, (!_, !_)) -> run interval0 l r Nothing -> pure ()@@ -383,8 +402,9 @@ | r' <= r = do -- contained onDel l' r' x'- IM.delete_ dim l'- res <- IM.lookupGT dim l'+ res <- stToPrim $ do+ IM.delete_ dim l'+ IM.lookupGT dim l' case res of Just rng -> run rng l r Nothing -> pure ()@@ -392,12 +412,13 @@ -- intersecting onDel l' r' x' onAdd r r' x'- IM.delete_ dim l'- IM.insert dim r (r', x')- pure ()+ stToPrim $ do+ IM.delete_ dim l'+ IM.insert dim r (r', x')+ pure () handleLeft l r = do- res <- IM.lookupLT dim l+ res <- stToPrim $ IM.lookupLT dim l case res of Nothing -> pure () Just (!l', (!r', !x'))@@ -409,19 +430,21 @@ onDel l' r' x' onAdd l' l x' onAdd r r' x'- -- IM.delete dim l'- IM.insert dim l' (l, x')- IM.insert dim r (r', x')+ stToPrim $ do+ -- IM.delete dim l'+ IM.insert dim l' (l, x')+ IM.insert dim r (r', x') | otherwise -> do -- intersecting onDel l' r' x' onAdd l' l x' -- IM.delete_ dim l'- IM.insert dim l' (l, x')+ stToPrim $ IM.insert dim l' (l, x') -- | \(O(\log n)\) Shorthand for overwriting the value of an interval that contains \([l, r)\). -- -- @since 1.1.0.0+{-# INLINE overwrite #-} overwrite :: (PrimMonad m, Eq a, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> a -> m () overwrite itm l r x = do res <- lookup itm l r@@ -433,6 +456,7 @@ -- Tracks interval state changes via @onAdd@ and @onDel@ hooks. -- -- @since 1.1.0.0+{-# INLINE overwriteM #-} overwriteM :: (PrimMonad m, Eq a, VU.Unbox a) => -- | The map@@ -458,5 +482,6 @@ -- where \([l, r)\) is the interval and \(x\) is the associated value. -- -- @since 1.1.0.0+{-# INLINE freeze #-} freeze :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> m (VU.Vector (Int, (Int, a))) freeze = IM.assocs . unITM
src/AtCoder/Extra/MultiSet.hs view
@@ -117,6 +117,7 @@ -- | \(O(n)\) Creates a `MultiSet` with capacity \(n\). -- -- @since 1.1.0.0+{-# INLINE new #-} new :: (PrimMonad m) => Int -> m (MultiSet (PrimState m)) new n = do mapMS <- HM.new n@@ -127,12 +128,14 @@ -- hash map. -- -- @since 1.1.0.0+{-# INLINE capacity #-} capacity :: MultiSet s -> Int capacity = HM.capacity . mapMS -- | \(O(1)\) Returns the number of distinct keys with positive counts. -- -- @since 1.1.0.0+{-# INLINE size #-} size :: (PrimMonad m) => MultiSet (PrimState m) -> m Int size MultiSet {..} = do VGM.unsafeRead cntMS 0@@ -140,6 +143,7 @@ -- | \(O(1)\) Looks up the count for a key. -- -- @since 1.1.0.0+{-# INLINE lookup #-} lookup :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m (Maybe Int) lookup MultiSet {..} k = do HM.lookup mapMS k <&> \case@@ -149,6 +153,7 @@ -- | \(O(1)\) Tests whether \(k\) is in the set. -- -- @since 1.1.0.0+{-# INLINE member #-} member :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m Bool member MultiSet {..} k = do HM.lookup mapMS k <&> \case@@ -158,18 +163,21 @@ -- | \(O(1)\) Tests whether \(k\) is not in the set. -- -- @since 1.1.0.0+{-# INLINE notMember #-} notMember :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m Bool notMember ms k = not <$> member ms k -- | \(O(1)\) Increments the count of a key. -- -- @since 1.1.0.0+{-# INLINE inc #-} inc :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m () inc ms k = add ms k 1 -- | \(O(1)\) Decrements the count of a key. -- -- @since 1.1.0.0+{-# INLINE dec #-} dec :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m () dec ms k = sub ms k 1 @@ -177,6 +185,7 @@ -- the \((k, c)\) pair is inserted. If \(v\) is negative, it falls back to `sub`. -- -- @since 1.1.0.0+{-# INLINE add #-} add :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m () add ms@MultiSet {..} k v = case compare v 0 of LT -> sub ms k (-v)@@ -195,6 +204,7 @@ -- `add`. -- -- @since 1.1.0.0+{-# INLINE sub #-} sub :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m () sub ms@MultiSet {..} k v = case compare v 0 of LT -> add ms k (-v)@@ -214,6 +224,7 @@ -- | \(O(1)\) Inserts a key-count pair into the set. `MultiSet` is actually a count map. -- -- @since 1.1.0.0+{-# INLINE insert #-} insert :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m () insert MultiSet {..} k v | v <= 0 = error $ "AtCoder.Extra.Multiset.insert: new count must be positive`" ++ show k ++ "`: `" ++ show v ++ "`"@@ -229,6 +240,7 @@ -- number of distinct keys that can be inserted into the internal hash map. -- -- @since 1.1.0.0+{-# INLINE delete #-} delete :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m () delete MultiSet {..} k = do HM.lookup mapMS k >>= \case
src/AtCoder/Extra/Pdsu.hs view
@@ -46,7 +46,8 @@ import AtCoder.Internal.Assert qualified as ACIA import Control.Monad-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Vector qualified as V import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM@@ -130,9 +131,14 @@ -- @since 1.1.0.0 {-# INLINE leader #-} leader :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> m Int-leader Pdsu {..} v0 = inner v0+leader pdsu v0 = stToPrim $ leaderST pdsu v0 where- !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.leader" v0 nPdsu+ !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.leader" v0 $ nPdsu pdsu++{-# INLINE leaderST #-}+leaderST :: (Semigroup a, VU.Unbox a) => Pdsu s a -> Int -> ST s Int+leaderST Pdsu {..} v0 = inner v0+ where inner v = do p <- VGM.read parentOrSizePdsu v if {- size? -} p < 0@@ -156,9 +162,9 @@ -- @since 1.1.0.0 {-# INLINE pot #-} pot :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> m a-pot dsu@Pdsu {..} v1 = do+pot dsu@Pdsu {..} v1 = stToPrim $ do -- Perform path compression- _ <- leader dsu v1+ _ <- leaderST dsu v1 VGM.read potentialPdsu v1 where !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.pot" v1 nPdsu@@ -169,7 +175,13 @@ -- @since 1.1.0.0 {-# INLINE same #-} same :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> m Bool-same !dsu !v1 !v2 = (==) <$> leader dsu v1 <*> leader dsu v2+same !dsu !v1 !v2 = stToPrim $ do+ l1 <- leaderST dsu v1+ l2 <- leaderST dsu v2+ pure $ l1 == l2+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.same" v1 $ nPdsu dsu+ !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.same" v2 $ nPdsu dsu -- TODO: call it unsafeDiff @@ -204,13 +216,18 @@ -- @since 1.1.0.0 {-# INLINE merge #-} merge :: (HasCallStack, PrimMonad m, Monoid a, Ord a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m Bool-merge dsu@Pdsu {..} v10 v20 !dp0 = inner v10 v20 dp0+merge dsu v10 v20 !dp0 = stToPrim $ mergeST dsu v10 v20 dp0 where- !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v10 nPdsu- !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v20 nPdsu+ !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v10 $ nPdsu dsu+ !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v20 $ nPdsu dsu++{-# INLINE mergeST #-}+mergeST :: (HasCallStack, Monoid a, Ord a, VU.Unbox a) => Pdsu s a -> Int -> Int -> a -> ST s Bool+mergeST dsu@Pdsu {..} v10 v20 !dp0 = inner v10 v20 dp0+ where inner v1 v2 !dp = do- !r1 <- leader dsu v1- !r2 <- leader dsu v2+ !r1 <- leaderST dsu v1+ !r2 <- leaderST dsu v2 if r1 == r2 then pure False else do@@ -274,17 +291,21 @@ -- @since 1.1.0.0 {-# INLINE size #-} size :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> m Int-size !dsu !v = (negate <$>) . VGM.read (parentOrSizePdsu dsu) =<< leader dsu v+size !dsu !v = stToPrim $ do+ l <- leaderST dsu v+ negate <$> VGM.read (parentOrSizePdsu dsu) l+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.size" v $ nPdsu dsu -- | \(O(n)\) Divides the graph into connected components and returns the list of them. -- -- @since 1.1.0.0 {-# INLINE groups #-} groups :: (PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> m (V.Vector (VU.Vector Int))-groups dsu@Pdsu {..} = do+groups dsu@Pdsu {..} = stToPrim $ do groupSize <- VUM.replicate nPdsu (0 :: Int) leaders <- VU.generateM nPdsu $ \i -> do- li <- leader dsu i+ li <- leaderST dsu i VGM.modify groupSize (+ 1) li pure li result <- do
src/AtCoder/Extra/Semigroup/Matrix.hs view
@@ -12,6 +12,7 @@ -- * Constructors new,+ square, zero, ident, diag,@@ -29,6 +30,17 @@ pow, powMod, powMint,++ -- * Rank+ rank,++ -- * Inverse+ inv,+ invRaw,++ -- * Determinant+ detMod,+ detMint, ) where @@ -36,13 +48,18 @@ import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Barrett qualified as BT import AtCoder.ModInt qualified as M+import Control.Monad (when)+import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.ST (runST) import Data.Foldable (for_) import Data.Semigroup (Semigroup (..)) import Data.Vector qualified as V import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Mutable qualified as VM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM+import Data.Word (Word64) import GHC.Exts (proxy#) import GHC.Stack (HasCallStack) import GHC.TypeNats (KnownNat, natVal')@@ -83,6 +100,13 @@ | VU.length vec /= h * w = error "AtCoder.Extra.Matrix: size mismatch" | otherwise = Matrix h w vec +-- | \(O(n^2)\) Creates an NxN square matrix.+--+-- @since 1.1.1.0+{-# INLINE square #-}+square :: (HasCallStack, VU.Unbox a) => Int -> VU.Vector a -> Matrix a+square n = new n n+ -- | \(O(n^2)\) Creates an NxN zero matrix. -- -- @since 1.1.0.0@@ -101,6 +125,8 @@ VGM.write vec (i + n * i) 1 pure vec +-- FIXME: diag should not take `n`+ -- | \(O(n^2)\) Creates an NxN diagonal matrix. -- -- @since 1.1.0.0@@ -134,26 +160,17 @@ -- -- @since 1.1.0.0 {-# INLINE mul #-}-mul :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e-mul !a !b =- Matrix h w' $- VU.unfoldrExactN- (h * w')- ( \(!row, !col) ->- let !x = f row col- in if col + 1 >= w'- then (x, (row + 1, 0))- else (x, (row, col + 1))- )- (0, 0)+mul :: forall e. (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.create $ do+ c <- VUM.replicate (h * w') (0 :: e)+ for_ [0 .. h - 1] $ \i -> do+ for_ [0 .. w - 1] $ \k -> do+ for_ [0 .. w' - 1] $ \j -> do+ let !aik = VG.unsafeIndex vecA (i * w + k)+ let !bkj = VG.unsafeIndex vecB (k * w' + j)+ VGM.unsafeModify c (+ (aik * bkj)) (i * w' + j)+ pure c where- f row col = VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (VU.unsafeSlice (w * row) w vecA)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch" -- | \(O(h_1 w_2 K)\) Multiplies H1xK matrix to a KxW2 matrix, taking the mod.@@ -161,7 +178,7 @@ -- @since 1.1.0.0 {-# INLINE mulMod #-} mulMod :: Int -> Matrix Int -> Matrix Int -> Matrix Int-mulMod !m !a !b =+mulMod !m (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -174,15 +191,14 @@ (0, 0) where !bt = BT.new32 $ fromIntegral m- f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod_ x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)- addMod x y = (x + y) `rem` m- mulMod_ x y = fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b+ -- NOTE: this is unsafe if the matrix is too large+ f row col =+ fromIntegral+ . (`rem` fromIntegral m)+ . VU.sum+ $ VU.imap+ (\iRow x -> BT.mulMod bt (fromIntegral x) (fromIntegral (VG.unsafeIndex vecB (col + (iRow * w')))))+ (VU.unsafeSlice (w * row) w vecA) !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMod: matrix size mismatch" -- | \(O(h_1 w_2 K)\) `mul` specialized to `M.ModInt`.@@ -196,7 +212,7 @@ {-# INLINE mulMintImpl #-} mulMintImpl :: forall a. (KnownNat a) => BT.Barrett -> Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)-mulMintImpl !bt !a !b =+mulMintImpl !bt (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.unfoldrExactN (h * w')@@ -208,31 +224,34 @@ ) (0, 0) where+ -- NOTE: this is unsafe if the matrix is too large f :: Int -> Int -> M.ModInt a- f row col = VU.sum $ VU.imap (\iRow x -> mulMod_ x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)- mulMod_ :: M.ModInt a -> M.ModInt a -> M.ModInt a- mulMod_ (M.ModInt x) (M.ModInt y) = M.unsafeNew . fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)- h = hM a- w = wM a- h' = hM b- vecA = vecM a- w' = wM b- vecB = vecM b- !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMint: matrix size mismatch"+ f row col =+ M.new64+ . VU.sum+ $ VU.imap+ ( \iRow x ->+ BT.mulMod+ bt+ (fromIntegral (M.unModInt x))+ (fromIntegral (M.unModInt (VG.unsafeIndex vecB (col + (iRow * w')))))+ )+ (VU.unsafeSlice (w * row) w vecA)+ !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMintImpl: matrix size mismatch" --- | \(O(w n^3)\) Calculates \(M^k\).+-- | \(O(w n^3)\) Returns \(M^k\). -- -- @since 1.1.0.0 {-# INLINE pow #-} pow :: Int -> Matrix Int -> Matrix Int pow k mat- | k < 0 = error "AtCoder.Extra.Matrix.powMod: the exponential must be non-negative"+ | k < 0 = error "AtCoder.Extra.Matrix.pow: the exponential must be non-negative" | k == 0 = ident $ hM mat | otherwise = ACEM.power mul k mat where- !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMod: matrix size mismatch"+ !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.pow: matrix size mismatch" --- | \(O(w n^3)\) Calculates \(M^k\), taking the mod.+-- | \(O(w n^3)\) Returns \(M^k\), taking the mod. -- -- @since 1.1.0.0 {-# INLINE powMod #-}@@ -244,9 +263,10 @@ where !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMod: matrix size mismatch" --- | \(O(w n^3)\) Calculates \(M^k\), specialized to `M.ModInt`.+-- | \(O(w n^3)\) Returns \(M^k\), specialized to `M.ModInt`. -- -- @since 1.1.0.0+{-# INLINE powMint #-} powMint :: forall m. (KnownNat m) => Int -> Matrix (M.ModInt m) -> Matrix (M.ModInt m) powMint k mat | k < 0 = error "AtCoder.Extra.Matrix.powMint: the exponential must be non-negative"@@ -256,9 +276,221 @@ !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMint: matrix size mismatch" !bt = BT.new32 $ fromIntegral (natVal' (proxy# @m)) +-- | (Internal)+{-# INLINE read2d #-}+read2d ::+ (PrimMonad m, VU.Unbox a) =>+ VM.MVector (PrimState m) (VUM.MVector (PrimState m) a) ->+ Int ->+ Int ->+ m a+read2d view i j = do+ row <- VGM.unsafeRead view i+ VGM.unsafeRead row j++-- | \(O(hw \min(h, w))\) Returns the rank of the matrix.+--+-- @since 1.1.1.0+{-# INLINE rank #-}+rank :: (Fractional a, Eq a, VU.Unbox a) => Matrix a -> Int+rank (Matrix h w vec) = runST $ do+ vm <- VU.thaw vec+ view <- V.thaw $ V.unfoldrExactN h (VUM.splitAt w) vm+ let inner rk j+ | rk == h || j == w = pure rk+ | otherwise = do+ xrj <- read2d view rk j+ when (xrj == 0) $ do+ let runSwap i+ | i == h = pure ()+ | otherwise = do+ xij <- read2d view i j+ if xij /= 0+ then VGM.unsafeSwap view rk i+ else runSwap (i + 1)+ runSwap (rk + 1)+ xrj' <- read2d view rk j+ if xrj' == 0+ then inner rk (j + 1)+ else do+ let c = 1 / xrj'+ rowRk <- VGM.read view rk+ -- for_ [j .. w - 1] $ \k -> do+ VGM.iforM_ (VGM.unsafeDrop j rowRk) $ \k_ x -> do+ VGM.unsafeWrite rowRk (k_ + j) $! c * x+ for_ [rk + 1 .. h - 1] $ \i -> do+ c <- read2d view i j+ rowI <- VGM.read view i+ -- for_ [j .. w - 1] $ \k -> do+ VGM.iforM_ (VGM.unsafeDrop j rowRk) $ \k_ ark -> do+ VGM.unsafeModify rowI (subtract (ark * c)) (k_ + j)+ inner (rk + 1) (j + 1)+ inner 0 0++-- TODO: add HasCallStack and compare their speeds++-- | \(O(n^3)\) Returns @(det, invMatrix)@ or `Nothing` if the matrix does not have inverse (the+-- determinant is zero).+--+-- ==== Constraints+-- - The input must be a square matrix.+--+-- @since 1.1.1.0+{-# INLINE inv #-}+inv :: forall a. (Fractional a, Eq a, VU.Unbox a) => Matrix a -> Maybe (a, Matrix a)+inv mat@(Matrix n _ _) = do+ (!det, !invMat) <- invRaw mat+ let !invMat' = VU.concat $ V.toList invMat+ pure (det, Matrix n n invMat')++-- | \(O(n^3)\) Returns @(det, invMatrix)@ or `Nothing` if the matrix does not have inverse (the+-- determinant is zero).+--+-- ==== Constraints+-- - The input must be a square matrix.+--+-- @since 1.1.1.0+{-# INLINE invRaw #-}+invRaw :: forall a. (Fractional a, Eq a, VU.Unbox a) => Matrix a -> Maybe (a, V.Vector (VU.Vector a))+invRaw (Matrix h w vec) = runST $ do+ vecA <- VU.thaw vec+ viewA <- V.thaw $ V.unfoldrExactN n (VUM.splitAt n) vecA++ vecB <- VUM.replicate (n * n) (0 :: a)+ for_ [0 .. n - 1] $ \i -> do+ VGM.unsafeWrite vecB (n * i + i) (1 :: a)+ viewB <- V.thaw $ V.unfoldrExactN n (VUM.splitAt n) vecB++ let inner i !det+ | i >= n = do+ viewB' <- V.mapM VU.unsafeFreeze =<< V.unsafeFreeze viewB+ pure $ Just (det, viewB')+ | otherwise = do+ let swapLoop k !det+ | k >= n = pure det+ | otherwise = do+ aki <- read2d viewA k i+ if aki /= 0+ then do+ if k /= i+ then do+ VGM.unsafeSwap viewA i k+ VGM.unsafeSwap viewB i k+ pure (-det)+ else pure det+ else do+ swapLoop (k + 1) det+ det' <- swapLoop i det+ aii <- read2d viewA i i+ if aii == 0+ then pure Nothing+ else do+ let !c = (1 :: a) / aii+ let !det'' = det' * aii+ rowAI <- VGM.unsafeRead viewA i+ rowBI <- VGM.unsafeRead viewB i+ VUM.iforM_ (VUM.unsafeDrop i rowAI) $ \j_ x -> do+ VGM.unsafeWrite rowAI (j_ + i) $! x * c+ VUM.iforM_ rowBI $ \j x -> do+ VGM.unsafeWrite rowBI j $! x * c+ for_ [0 .. n - 1] $ \k -> do+ when (i /= k) $ do+ c <- read2d viewA k i+ rowAK <- VGM.unsafeRead viewA k+ rowBK <- VGM.unsafeRead viewB k+ VGM.iforM_ (VGM.unsafeDrop i rowAI) $ \j_ aij -> do+ VGM.unsafeModify rowAK (subtract (aij * c)) (j_ + i)+ VGM.iforM_ rowBI $ \j bij -> do+ VGM.unsafeModify rowBK (subtract (bij * c)) j+ inner (i + 1) det''++ inner 0 (1 :: a)+ where+ !_ = ACIA.runtimeAssert (h == w) $ "AtCoder.Extra.Semigroup.Matrix.inv: given non-square matrix of size " ++ show (h, w)+ n = h++-- | \(O(hw \min(h, w))\) Returns the rank of the matrix.+--+-- @since 1.1.1.0+{-# INLINE detMod #-}+detMod :: Int -> Matrix Int -> Int+detMod m (Matrix h w vecA) = runST $ do+ vm <- VU.thaw vecA+ view <- V.thaw $ V.unfoldrExactN n (VUM.splitAt n) vm++ let inner i (!det :: Int)+ | i >= n = pure det+ | otherwise = do+ let swapLoop j !det+ | j >= n = pure det+ | otherwise = do+ aji <- read2d view j i+ if aji == 0+ then swapLoop (j + 1) det+ else do+ if i /= j+ then do+ VGM.unsafeSwap view i j+ pure $! m - det+ else pure det+ det' <- swapLoop i det+ det'' <- VU.foldM'+ ( \ !acc j -> do+ let visitDiag !det = do+ aii <- read2d view i i+ if aii == 0+ then pure det+ else do+ aji <- read2d view j i+ let !c = m - aji `div` aii+ rowI <- VGM.unsafeRead view i+ rowJ <- VGM.unsafeRead view j+ -- NOTE: it's a reverse loop!+ VGM.ifoldrM'+ ( \k_ aik () -> do+ VGM.unsafeModify rowJ ((`mod` m) . (+ aik * c)) (k_ + i)+ )+ ()+ (VGM.unsafeDrop i rowI)+ VGM.unsafeSwap view i j+ visitDiag (m - det)+ acc' <- visitDiag acc+ VGM.unsafeSwap view i j+ pure $! m - acc'+ )+ det'+ (VU.generate (n - (i + 1)) (+ (i + 1)))++ inner (i + 1) det''++ det <- inner 0 (1 :: Int)+ fromIntegral+ <$> VU.foldM'+ ( \(!acc :: Word64) i -> do+ aii <- read2d view i i+ pure $! BT.mulMod bt acc $! fromIntegral aii+ )+ (fromIntegral det)+ (VU.generate n id)+ where+ !_ = ACIA.runtimeAssert (h == w) $ "AtCoder.Extra.Semigroup.Matrix.detMod: given non-square matrix of size " ++ show (h, w)+ !n = h+ !bt = BT.new32 $ fromIntegral m++-- | \(O(hw \min(h, w))\) Returns the rank of the matrix.+--+-- @since 1.1.1.0+{-# INLINE detMint #-}+detMint :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> M.ModInt a+detMint matA = M.new . detMod m $ map M.val matA+ where+ !m = fromIntegral (natVal' (proxy# @a))+ -- | @since 1.1.0.0 instance (Num a, VU.Unbox a) => Semigroup (Matrix a) where {-# INLINE (<>) #-} (<>) = mul++ -- Prefer `powMod` or `powMint` as specialized, much efficient variant. {-# INLINE stimes #-} stimes = ACEM.power (<>) . fromIntegral
src/AtCoder/Extra/Tree/Hld.hs view
@@ -216,14 +216,14 @@ -- | \(O(n)\) Creates an `Hld` with \(0\) as the root vertex. -- -- @since 1.1.0.0-{-# INLINE new #-}+{-# INLINABLE new #-} new :: forall w. (HasCallStack) => Gr.Csr w -> Hld new tree = newAt tree 0 -- | \(O(n)\) Creates an `Hld` with a root vertex specified. -- -- @since 1.1.0.0-{-# INLINE newAt #-}+{-# INLINABLE newAt #-} newAt :: forall w. (HasCallStack) => Gr.Csr w -> Vertex -> Hld newAt tree root = runST $ do -- Re-create adjacent vertices so that the biggest subtree's head vertex comes first.@@ -315,7 +315,7 @@ -- | \(O(\log n)\) Calculates the lowest common ancestor of \(u\) and \(v\). -- -- @since 1.1.0.0-{-# INLINE lca #-}+{-# INLINABLE lca #-} lca :: (HasCallStack) => Hld -> Vertex -> Vertex -> Vertex lca Hld {..} = inner where@@ -338,7 +338,7 @@ -- is bigger than the depth of \(v\). -- -- @since 1.1.0.0-{-# INLINE ancestor #-}+{-# INLINABLE ancestor #-} ancestor :: (HasCallStack) => Hld -> Vertex -> Int -> Vertex ancestor Hld {..} parent k0 = inner parent k0 where@@ -357,7 +357,7 @@ -- Throws an error if `k` is out -- -- @since 1.1.0.0-{-# INLINE jump #-}+{-# INLINABLE jump #-} jump :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int -> Maybe Vertex jump hld@Hld {..} u v k | k > lenU + lenV = Nothing@@ -373,7 +373,7 @@ -- | \(O(\log n)\) Returns the length of the path between \(u\) and \(v\). -- -- @since 1.1.0.0-{-# INLINE lengthBetween #-}+{-# INLINABLE lengthBetween #-} lengthBetween :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int lengthBetween hld@Hld {..} u v = du - dLca + dv - dLca where@@ -385,7 +385,7 @@ -- | \(O(n)\) Returns the vertices on the path between \(u\) and \(v\). -- -- @since 1.1.0.0-{-# INLINE path #-}+{-# INLINABLE path #-} path :: (HasCallStack) => Hld -> Vertex -> Vertex -> [Vertex] path hld@Hld {..} u v = concatMap expand $ pathSegmentsInclusive WeightsAreOnVertices hld u v where@@ -400,7 +400,7 @@ -- `WeightsAreOnEdges`. This is the trick to put edge weights to on vertices. -- -- @since 1.1.0.0-{-# INLINE pathSegmentsInclusive #-}+{-# INLINABLE pathSegmentsInclusive #-} pathSegmentsInclusive :: (HasCallStack) => WeightPolicy -> Hld -> Vertex -> Vertex -> [(VertexHld, VertexHld)] pathSegmentsInclusive weightPolicy Hld {..} x0 y0 = done $ inner x0 [] y0 [] where@@ -440,7 +440,7 @@ -- corresponds to the subtree segments rooted at the given @subtreeRoot@. -- -- @since 1.1.0.0-{-# INLINE subtreeSegmentInclusive #-}+{-# INLINABLE subtreeSegmentInclusive #-} subtreeSegmentInclusive :: (HasCallStack) => Hld -> Vertex -> (VertexHld, VertexHld) subtreeSegmentInclusive Hld {..} subtreeRoot = (ir, ir + sr - 1) where@@ -450,7 +450,7 @@ -- | \(O(1)\) Returns `True` if \(u\) is in a subtree of \(r\). -- -- @since 1.1.0.0-{-# INLINE isInSubtree #-}+{-# INLINABLE isInSubtree #-} isInSubtree :: (HasCallStack) => Hld -> Vertex -> Vertex -> Bool isInSubtree hld@Hld {..} r_ u = l <= iu && iu <= r where
+ src/AtCoder/Extra/Tree/Lct.hs view
@@ -0,0 +1,844 @@+{-# LANGUAGE DerivingVia #-}+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE TypeFamilies #-}++-- | Link/cut tree: forest with monoid values.+--+-- ==== __Example__+--+-- Create a link/cut tree of @Sum Int@ with inverse operator `negate`:+--+-- >>> import AtCoder.Extra.Tree.Lct qualified as Lct+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> -- 0--1--2+-- >>> -- +--3+-- >>> lct <- Lct.buildInv negate (VU.generate 4 Sum) $ VU.fromList [(0, 1), (1, 2), (1, 3)]+--+-- Monoid products can be calculated for paths or subtrees:+--+-- >>> Lct.prodPath lct 0 2+-- Sum {getSum = 3}+--+-- >>> Lct.prodSubtree lct 1 {- parent -} 2+-- Sum {getSum = 4}+--+-- `root` returns the current root vertex of the underlying tree, which is not easy to predict:+--+-- >>> Lct.root lct 3+-- 2+--+-- Set (`evert`) the root of the underlying tree to \(0\) and get the `lca` of vertices \(2\) and+-- \(3\):+--+-- >>> Lct.evert lct 0+-- >>> Lct.lca lct 2 3+-- 1+--+-- Similar to @Hld@, `Lct` allows various tree queries:+--+-- >>> Lct.parent lct 3+-- Just 1+--+-- >>> Lct.jump lct 2 3 2+-- 3+--+-- Edges can be dynamically added (`link`) or removed (`cut`):+--+-- >>> -- 0 1 2+-- >>> -- +--3+-- >>> Lct.cut lct 0 1+-- >>> Lct.cut lct 1 2+-- >>> VU.generateM 4 (Lct.root lct)+-- [0,1,2,1]+--+-- >>> -- +-----++-- >>> -- 0 1 2+-- >>> -- +--3+-- >>> Lct.link lct 0 2+-- >>> VU.generateM 4 (Lct.root lct)+-- [2,1,2,1]+--+-- @since 1.1.1.0+module AtCoder.Extra.Tree.Lct+ ( -- Link/cut tree+ Lct (..),+ Vertex,++ -- * Constructors+ new,+ newInv,+ build,+ buildInv,++ -- * Modifications++ -- ** Write+ write,+ modify,+ modifyM,++ -- ** Link/cut+ link,+ cut,++ -- ** Evert/expose+ evert,+ expose,+ expose_,++ -- * Tree queries++ -- ** Root, parent, jump, LCA+ root,+ parent,+ jump,+ lca,++ -- ** Products+ prodPath,+ prodSubtree,+ )+where++import AtCoder.Internal.Assert qualified as ACIA+import Control.Monad (unless, when)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST)+import Data.Bit+import Data.Bits+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)++-- import GHC.Stack (HasCallStack)++-- | Alias of vertex type.+type Vertex = Int++{-# INLINE undefLct #-}+undefLct :: Vertex+undefLct = -1++{-# INLINE nullLct #-}+nullLct :: Vertex -> Bool+nullLct = (== -1)++-- We could optimize the with options, but++-- | Link/cut tree.+--+-- @since 1.1.1.0+data Lct s a = Lct+ { -- | The number of vertices.+ --+-- @since 1.1.1.0+ nLct :: {-# UNPACK #-} !Int,+ -- | Decomposed node data storage: left children.+ --+ -- @since 1.1.1.0+ lLct :: !(VUM.MVector s Vertex),+ -- | Decomposed node data storage: right children.+ --+ -- @since 1.1.1.0+ rLct :: !(VUM.MVector s Vertex),+ -- | Decomposed node data storage: parents.+ --+ -- @since 1.1.1.0+ pLct :: !(VUM.MVector s Vertex),+ -- | Decomposed node data storage: subtree sizes.+ --+ -- @since 1.1.1.0+ sLct :: !(VUM.MVector s Int),+ -- | Decomposed node data storage: reverse flag.+ --+ -- @since 1.1.1.0+ revLct :: !(VUM.MVector s Bit),+ -- | Decomposed node data storage: monoid values.+ --+ -- @since 1.1.1.0+ vLct :: !(VUM.MVector s a),+ -- | Decomposed node data storage: monoid products.+ --+ -- @since 1.1.1.0+ prodLct :: !(VUM.MVector s a),+ -- | Decomposed node data storage: dual monod product (right fold). This is required for+ -- non-commutative monoids only.+ --+ -- @since 1.1.1.0+ dualProdLct :: !(VUM.MVector s a),+ -- | Decomposed node data storage: path-parent monoid product. This works for subtree product+ -- queries over commutative monoids only.+ --+ -- @since 1.1.1.0+ midLct :: !(VUM.MVector s a),+ -- | Decomposed node data storage: monoid product of subtree. This works for subtree product+ -- queries over commutative monoids only.+ --+ -- @since 1.1.1.0+ subtreeProdLct :: !(VUM.MVector s a),+ -- | Inverse operator of the monoid. This works for subtree product queries over commutative+ -- monoids only.+ --+ -- @since 1.1.1.0+ invOpLct :: !(a -> a)+ }++-- | \(O(n)\) Creates a link/cut tree with \(n\) vertices and no edges.+--+-- @since 1.1.1.0+{-# INLINE new #-}+new :: (PrimMonad m, Monoid a, VU.Unbox a) => Int -> m (Lct (PrimState m) a)+new = newInv id++-- | \(O(n + m \log n)\) Creates a link/cut tree with an inverse operator, initial monoid values and+-- no edges. This setup enables subtree queries (`prodSubtree`).+--+-- @since 1.1.1.0+{-# INLINE newInv #-}+newInv :: (PrimMonad m, Monoid a, VU.Unbox a) => (a -> a) -> Int -> m (Lct (PrimState m) a)+newInv !invOpLct nLct = buildInv invOpLct (VU.replicate nLct mempty) VU.empty++-- | \(O(n + m \log n)\) Creates a link/cut tree of initial monoid values and initial edges.+--+-- @since 1.1.1.0+{-# INLINE build #-}+build ::+ (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+ -- | Vertex monoid values+ VU.Vector a ->+ -- | Edges+ VU.Vector (Vertex, Vertex) ->+ -- | Link/cut tree+ m (Lct (PrimState m) a)+build = buildInv id++-- | \(O(n + m \log n)\) Creates a link/cut tree with an inverse operator, initial monoid values and+-- initial edges. This setup enables subtree queries (`prodSubtree`).+--+-- @since 1.1.1.0+{-# INLINE buildInv #-}+buildInv ::+ (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+ -- | Inverse operator+ (a -> a) ->+ -- | Vertex monoid values+ VU.Vector a ->+ -- | Edges+ VU.Vector (Vertex, Vertex) ->+ -- | Link/cut tree+ m (Lct (PrimState m) a)+buildInv !invOpLct xs es = do+ lct <- do+ let !nLct = VU.length xs+ lLct <- VUM.replicate nLct undefLct+ rLct <- VUM.replicate nLct undefLct+ pLct <- VUM.replicate nLct undefLct+ sLct <- VUM.replicate nLct 0+ revLct <- VUM.replicate nLct (Bit False)+ vLct <- VU.thaw xs+ prodLct <- VUM.replicate nLct mempty+ dualProdLct <- VUM.replicate nLct mempty+ midLct <- VUM.replicate nLct mempty+ subtreeProdLct <- VUM.replicate nLct mempty+ pure Lct {..}+ VU.forM_ es $ \(!u, !v) -> do+ link lct u v+ pure lct++-- -------------------------------------------------------------------------------------------------+-- Balancing+-- -------------------------------------------------------------------------------------------------++-- | \(O(1)\) Rotates up a non-root node.+{-# INLINEABLE rotateST #-}+rotateST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s ()+rotateST lct@Lct {pLct, lLct, rLct} v = do+ p <- VGM.unsafeRead pLct v+ pp <- VGM.unsafeRead pLct p+ pl <- VGM.unsafeRead lLct p++ c <-+ if pl == v+ then do+ -- rotate right:+ -- p v <-- reference from `pp` is updated later+ -- / \+ -- v -> p+ -- \ /+ -- c c+ c <- VGM.unsafeExchange rLct v p+ VGM.unsafeWrite lLct p c+ pure c+ else do+ -- rotate left:+ -- p v <-- reference from `pp` is updated later+ -- \ /+ -- v -> p+ -- / \+ -- c c+ c <- VGM.unsafeExchange lLct v p+ VGM.unsafeWrite rLct p c+ pure c++ updateNodeST lct p+ updateNodeST lct v++ -- update the reference from `pp`:+ unless (nullLct pp) $ do+ ppl <- VGM.unsafeRead lLct pp+ if ppl == p+ then VGM.unsafeWrite lLct pp v+ else do+ ppr <- VGM.unsafeRead rLct pp+ if ppr == p+ then VGM.unsafeWrite rLct pp v+ else do+ -- overwrite the light (path-parent) pointer:+ changeLightST lct pp p v++ -- update parent pointers to `pp`: pp <-- v <-- p <-- c+ VGM.unsafeWrite pLct v pp+ VGM.unsafeWrite pLct p v+ unless (nullLct c) $ do+ VGM.unsafeWrite pLct c p++-- | Amortized \(O(\log n)\). Moves a node up to the root, performing self-balancing heuristic+-- called rotations.+{-# INLINEABLE splayST #-}+splayST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s ()+splayST lct@Lct {pLct} c = do+ pushNodeST lct c+ let inner = do+ isRootC <- isRootNodeST lct c+ unless isRootC $ do+ p <- VGM.unsafeRead pLct c+ pp <- if nullLct p then pure undefLct else VGM.unsafeRead pLct p+ placeP <- nodePlaceST lct p+ if placeP == RootNodeLct+ then do+ pushNodeST lct p+ pushNodeST lct c+ rotateST lct c+ else do+ placeC <- nodePlaceST lct c+ pushNodeST lct pp+ pushNodeST lct p+ pushNodeST lct c+ if placeC == placeP+ then do+ -- Rotate right twice:+ --+ -- pp p c+ -- / / \ \+ -- p -> c pp -> p+ -- / \+ -- c pp++ -- Or rotate left twice:+ --+ -- pp p c+ -- \ / \ /+ -- p -> pp c -> p+ -- \ /+ -- c pp++ rotateST lct p+ rotateST lct c+ else do+ -- pp pp c+ -- / / | \+ -- p -> c -> p pp+ -- \ /+ -- c p+ rotateST lct c+ rotateST lct c+ inner+ inner++-- * Node helpers++-- | \(O(1)\)+{-# INLINEABLE isRootNodeST #-}+isRootNodeST :: Lct s a -> Vertex -> ST s Bool+isRootNodeST lct v = do+ (== RootNodeLct) <$> nodePlaceST lct v++-- TODO: return heavy/light notion+data NodePlaceLct = RootNodeLct | LeftNodeLct | RightNodeLct+ deriving (Eq)++-- | \(O(1)\)+{-# INLINEABLE nodePlaceST #-}+nodePlaceST :: Lct s a -> Vertex -> ST s NodePlaceLct+nodePlaceST Lct {lLct, rLct, pLct} v = do+ p <- VGM.unsafeRead pLct v+ if nullLct p+ then pure RootNodeLct+ else do+ pl <- VGM.unsafeRead lLct p+ if pl == v+ then pure LeftNodeLct+ else do+ pr <- VGM.unsafeRead rLct p+ if pr == v+ then pure RightNodeLct+ else pure RootNodeLct++-- -------------------------------------------------------------------------------------------------+-- Node operations+-- -------------------------------------------------------------------------------------------------++-- | \(O(1)\) Propgates the lazily propagated values on a node.+{-# INLINEABLE pushNodeST #-}+pushNodeST :: (VU.Unbox a) => Lct s a -> Vertex -> ST s ()+pushNodeST lct@Lct {lLct, rLct, revLct} v = do+ Bit b <- VGM.unsafeExchange revLct v (Bit False)+ when b $ do+ l <- VGM.unsafeRead lLct v+ r <- VGM.unsafeRead rLct v+ unless (nullLct l) $ reverseNodeST lct l+ unless (nullLct r) $ reverseNodeST lct r++-- | \(O(1)\)+{-# INLINEABLE reverseNodeST #-}+reverseNodeST :: (VU.Unbox a) => Lct s a -> Vertex -> ST s ()+reverseNodeST lct@Lct {revLct} i = do+ -- lazily propagate new reverse from the children, or cancel:+ VGM.unsafeModify revLct (xor (Bit True)) i+ -- swap+ swapLrNodeST lct i++-- | \(O(1)\) Reverses the left and the right children, lazily and recursively.+{-# INLINEABLE swapLrNodeST #-}+swapLrNodeST :: (VU.Unbox a) => Lct s a -> Vertex -> ST s ()+swapLrNodeST Lct {lLct, rLct, prodLct, dualProdLct} i = do+ -- swap chidlren+ VGM.unsafeModifyM lLct (VGM.unsafeExchange rLct i) i+ -- swap prodLct[i] and dualProdLct[i]+ VGM.unsafeModifyM prodLct (VGM.unsafeExchange dualProdLct i) i++-- | \(O(1)\) Recomputes the node size and the monoid product.+{-# INLINEABLE updateNodeST #-}+updateNodeST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s ()+updateNodeST Lct {..} i = do+ l <- VGM.unsafeRead lLct i+ r <- VGM.unsafeRead rLct i+ v <- VGM.unsafeRead vLct i+ m <- VGM.unsafeRead midLct i++ (!size', !prod', !dualProd', !subtreeProd') <-+ if nullLct l+ then pure (1 :: Int, v, v, v <> m)+ else do+ lSize <- VGM.unsafeRead sLct l+ lProd <- VGM.unsafeRead prodLct l+ lDualProd <- VGM.unsafeRead dualProdLct l+ lSubtreeProd <- VGM.unsafeRead subtreeProdLct l+ pure (lSize + 1, lProd <> v, v <> lDualProd, lSubtreeProd <> v <> m)++ (!size'', !prod'', !dualProd'', !subtreeProd'') <-+ if nullLct r+ then pure (size', prod', dualProd', subtreeProd')+ else do+ rSize <- VGM.unsafeRead sLct r+ rProd <- VGM.unsafeRead prodLct r+ rDualProd <- VGM.unsafeRead dualProdLct r+ rSubtreeProd <- VGM.unsafeRead subtreeProdLct r+ pure (size' + rSize, prod' <> rProd, rDualProd <> dualProd', subtreeProd' <> rSubtreeProd)++ VGM.unsafeWrite sLct i size''+ VGM.unsafeWrite prodLct i prod''+ VGM.unsafeWrite dualProdLct i dualProd''+ VGM.unsafeWrite subtreeProdLct i subtreeProd''++-- | \(O(1)\) Called on adding a path-parent edge. This is for subtree folding.+{-# INLINEABLE addLightST #-}+addLightST :: (Semigroup a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s ()+addLightST Lct {subtreeProdLct, midLct} p c = do+ newChild <- VGM.unsafeRead subtreeProdLct c+ VGM.unsafeModify midLct (newChild <>) p++-- | \(O(1)\) Called on changing a path-parent edge. This is for subtree folding.+{-# INLINEABLE changeLightST #-}+changeLightST :: Lct s a -> Vertex -> Vertex -> Vertex -> ST s ()+changeLightST lct u v p = do+ pure ()++-- | \(O(1)\) Called on erasing a path-parent edge. This is for subtree folding.+{-# INLINEABLE eraseLightST #-}+eraseLightST :: (Semigroup a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s ()+eraseLightST Lct {subtreeProdLct, midLct, invOpLct} p c = do+ sub <- VGM.unsafeRead subtreeProdLct c+ let !sub' = invOpLct sub+ VGM.unsafeModify midLct (<> sub') p++-- -------------------------------------------------------------------------------------------------+-- Write+-- -------------------------------------------------------------------------------------------------++-- TODO: read++-- | Amortized \(O(\log n)\). Writes the monoid value of a vertex.+--+-- @since 1.1.1.0+{-# INLINE write #-}+write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> a -> m ()+write lct v x = stToPrim $ do+ -- make @v@ the new root of the underlying tree:+ evertST lct v+ VGM.unsafeWrite (vLct lct) v x+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.write" v (nLct lct)++-- | Amortized \(O(\log n)\). Modifies the monoid value of a vertex with a pure function.+--+-- @since 1.1.1.0+{-# INLINE modify #-}+modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> (a -> a) -> Vertex -> m ()+modify lct f v = stToPrim $ do+ -- make @v@ the new root of the underlying tree:+ evertST lct v+ VGM.unsafeModify (vLct lct) f v+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.modify" v (nLct lct)++-- | Amortized \(O(\log n)\). Modifies the monoid value of a vertex with a monadic function.+--+-- @since 1.1.1.0+{-# INLINE modifyM #-}+modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> (a -> m a) -> Vertex -> m ()+modifyM lct f v = do+ -- make @v@ the new root of the underlying tree:+ stToPrim $ evertST lct v+ VGM.unsafeModifyM (vLct lct) f v+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.modifyM" v (nLct lct)++-- -------------------------------------------------------------------------------------------------+-- Link/cut operations+-- -------------------------------------------------------------------------------------------------++-- | Amortized \(O(\log n)\).+{-# INLINEABLE linkST #-}+linkST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s ()+linkST lct@Lct {pLct, rLct} c p = do+ -- make @c@ the new root of the underlying tree+ evertST lct c+ -- remove right children of @p@.+ _ <- exposeST lct p+ pushNodeST lct p++ -- dbgM $ do+ -- cp <- VGM.unsafeRead pLct c+ -- let !_ = ACIA.runtimeAssert (nullLct cp) $ "cp must be null: " ++ show (c, cp)+ -- pr <- VGM.unsafeRead rLct p+ -- let !_ = ACIA.runtimeAssert (nullLct pr) $ "pr must be null: " ++ show (p, pr)+ -- pure ()++ -- connect with a heavy edge:+ VGM.unsafeWrite pLct c p+ VGM.unsafeWrite rLct p c+ updateNodeST lct p++-- | Amortized \(O(\log n)\). Creates an edge between \(c\) and \(p\). In the represented tree, the+-- parent of \(c\) will be \(p\) after this operation.+--+-- @since 1.1.1.0+{-# INLINE link #-}+link :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m ()+link lct c p = stToPrim $ linkST lct c p+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.link" c (nLct lct)+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.link" p (nLct lct)++{-# INLINEABLE cutST #-}+cutST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s ()+cutST lct@Lct {pLct, lLct} u v = do+ -- make @u@ the new root of the underlying tree+ evertST lct u+ -- make @v@ in the same preferred path as the root+ _ <- exposeST lct v++ -- dbgM $ do+ -- -- @v@ does not have any right children. because @u@ and @v@ are neighbors, @vl@ is @u@.+ -- vp <- VGM.unsafeRead pLct v+ -- let !_ = ACIA.runtimeAssert (nullLct vp) "vp must be null"+ -- vl <- VGM.unsafeRead lLct v+ -- let !_ = ACIA.runtimeAssert (vl == u) "vl must be `u`"+ -- pure ()++ -- do+ -- -- @v@ does not have any right children. because @u@ and @v@ are neighbors, @vl@ is @u@.+ -- vp <- VGM.unsafeRead pLct v+ -- vl <- VGM.unsafeRead lLct v+ -- let !_ = if nullLct vp then () else error "vp must be null"+ -- let !_ = if vl == u then () else error "vl must be `u`"+ -- pure ()++ -- delete the heavy edge.+ -- vl <- VGM.unsafeRead lLct v+ -- VGM.unsafeWrite pLct vl undefLct+ VGM.unsafeWrite pLct u undefLct+ VGM.unsafeWrite lLct v undefLct+ updateNodeST lct v++-- | Amortized \(O(\log N)\). Deletes an edge between \(u\) and \(v\).+--+-- @since 1.1.1.0+{-# INLINE cut #-}+cut :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m ()+cut lct u v = stToPrim $ cutST lct u v+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.cut" u (nLct lct)+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.cut" v (nLct lct)++-- | Amortized \(O(\log n)\). Makes \(v\) a new root of the underlying tree.+{-# INLINEABLE evertST #-}+evertST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s ()+evertST lct v = do+ -- make @v@ be in the same preferred path as root. note that @v@ is at the root of the auxiliary tree.+ _ <- exposeST lct v+ -- reverse all the edges with respect to @v@: make @v@ a new root of the auxiliary tree.+ reverseNodeST lct v+ pushNodeST lct v++-- | Amortized \(O(\log n)\). Makes \(v\) a new root of the underlying tree.+--+-- @since 1.1.1.0+{-# INLINE evert #-}+evert :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m ()+evert lct v = stToPrim $ evertST lct v+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.evert" v (nLct lct)++{-# INLINEABLE exposeST #-}+exposeST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s Vertex+exposeST lct@Lct {pLct, rLct} v0 = do+ let inner v lastRoot+ | nullLct v = pure lastRoot+ | otherwise = do+ -- go up to the top of the auxiliary tree:+ splayST lct v++ -- make @lastRoot@ the right child of @v@:+ -- v v+ -- /|\ -> /|\+ -- | r | lastRoot <-- @v0@ (in the @lastRoot@) will be connected to the root+ -- lastRoot r+ r <- VGM.unsafeRead rLct v+ unless (nullLct r) $ addLightST lct v r+ unless (nullLct lastRoot) $ eraseLightST lct v lastRoot+ VGM.unsafeWrite rLct v lastRoot+ updateNodeST lct v++ -- go up to the next auxiliary tree:+ -- p+ -- |+ -- v+ -- \+ -- lastRoot+ vp <- VGM.unsafeRead pLct v+ inner vp v++ res <- inner v0 undefLct++ -- do+ -- -- FIXME: remove+ -- pRes <- VGM.unsafeRead pLct res+ -- unless (nullLct pRes) $ error $ "xxx must be null!!! " ++ show (res, pRes)++ splayST lct v0++ -- do+ -- -- FIXME: remove+ -- p <- VGM.unsafeRead pLct v0+ -- unless (nullLct p) $ error $ "must be null!!! " ++ show (res, v0, p)++ pure res++-- | Amortized \(O(\log n)\). Makes \(v\) and the root to be in the same preferred path (auxiliary+-- tree). After the opeartion, \(v\) will be the new root and all the children will be detached from+-- the preferred path.+--+-- @since 1.1.1.0+{-# INLINE expose #-}+expose :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m Vertex+expose lct v = stToPrim $ exposeST lct v+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.expose_" v (nLct lct)++-- | Amortized \(O(\log n)\). `expose` with the return value discarded.+--+-- @since 1.1.1.0+{-# INLINE expose_ #-}+expose_ :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m ()+expose_ lct v0 = stToPrim $ do+ _ <- exposeST lct v0+ pure ()+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.expose_" v0 (nLct lct)++-- -------------------------------------------------------------------------------------------------+-- Jumo, LCA+-- -------------------------------------------------------------------------------------------------++-- | \(O(\log n)\) Returns the root of the underlying tree. Two vertices in the same connected+-- component have the same root vertex.+--+-- @since 1.1.1.0+{-# INLINE root #-}+root :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Int -> m Vertex+root lct@Lct {lLct} c0 = stToPrim $ do+ _ <- exposeST lct c0+ pushNodeST lct c0+ let inner c = do+ cl <- VGM.unsafeRead lLct c+ if nullLct cl+ then pure c+ else do+ pushNodeST lct cl+ inner cl+ c' <- inner c0+ splayST lct c'+ pure c'+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.root" c0 (nLct lct)++-- | \(O(\log n)\) Returns the parent vertex in the underlying tree.+--+-- @since 1.1.1.0+{-# INLINE parent #-}+parent :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Int -> m (Maybe Vertex)+parent lct@Lct {lLct, rLct} x = stToPrim $ do+ _ <- exposeST lct x+ pushNodeST lct x+ xl <- VGM.unsafeRead lLct x+ if nullLct xl+ then pure Nothing+ else do+ pushNodeST lct xl+ let inner y = do+ yr <- VGM.unsafeRead rLct y+ if nullLct yr+ then pure y+ else do+ pushNodeST lct yr+ inner yr+ Just <$> inner xl+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.parent" x (nLct lct)++{-# INLINEABLE jumpST #-}+jumpST :: (HasCallStack, Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> Int -> ST s Vertex+jumpST lct@Lct {lLct, rLct, sLct} u0 v0 k0 = do+ -- make @v0@ a new root of the underlying tree+ evertST lct v0+ -- make @u0@ in the same preferred path as the root (@v0)+ _ <- exposeST lct u0++ do+ size <- VGM.unsafeRead sLct u0+ let !_ = ACIA.runtimeAssert (0 <= k0 && k0 < size) "invalid jump"+ pure ()++ let inner k u = do+ pushNodeST lct u+ -- TODO: what is happening?+ ur <- VGM.unsafeRead rLct u+ urSize <- if nullLct ur then pure 0 else VGM.unsafeRead sLct ur+ case compare k urSize of+ LT -> inner k ur+ EQ -> pure u+ GT -> do+ ul <- VGM.unsafeRead lLct u+ inner (k - (urSize + 1)) ul++ res <- inner k0 u0+ splayST lct res+ pure res++-- | \(O(\log n)\) Given a path between \(u\) and \(v\), returns the \(k\)-th vertex of the path.+--+-- ==== Constraints+-- - The \(k\)-th vertex must exist.+--+-- @since 1.1.1.0+{-# INLINE jump #-}+jump :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> Int -> m Vertex+jump lct u v k = stToPrim $ jumpST lct u v k++-- | \(O(\log n)\) Returns the LCA of \(u\) and \(v\). Because the root of the underlying changes+-- in almost every operation, one might want to use `evert` beforehand.+--+-- ==== Constraints+-- - \(u\) and \(v\) must be in the same connected component.+--+-- @since 1.1.1.0+{-# INLINE lca #-}+lca :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Int -> Int -> m Vertex+lca lct u v = stToPrim $ do+ ru <- root lct u+ rv <- root lct v+ let !_ = ACIA.runtimeAssert (ru == rv) $ "AtCoder.Extra.Lct.lca: given two vertices in different connected components " ++ show (u, v)+ _ <- exposeST lct u+ exposeST lct v++-- -------------------------------------------------------------------------------------------------+-- Monoid produ\t+-- -------------------------------------------------------------------------------------------------++-- | Amortized \(O(\log n)\). Folds a path between \(u\) and \(v\) (inclusive).+--+-- @since 1.1.1.0+{-# INLINE prodPath #-}+prodPath :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m a+prodPath lct@Lct {prodLct} u v = stToPrim $ do+ -- make @u@ the root of the underlying tree+ evertST lct u+ -- make @v@ in the same preferred path as @u@+ _ <- exposeST lct v+ -- now that @v@ is at the root of the auxiliary tree, its aggregation value is the path folding:+ VGM.unsafeRead prodLct v+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.prodPath" u (nLct lct)+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.prodPath" v (nLct lct)++-- | Amortized \(O(\log n)\). Fold the subtree under \(v\), considering \(p\) as the root-side+-- vertex. Or, if \(p\) equals to \(v\), \(v\) will be the new root.+--+-- ==== Constraints+-- - The inverse operator has to be set on consturction (`newInv` or `buildInv`).+--+-- @since 1.1.1.0+{-# INLINE prodSubtree #-}+prodSubtree ::+ (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+ -- | Link/cut tree+ Lct (PrimState m) a ->+ -- | Vertex+ Vertex ->+ -- | Root or parent+ Vertex ->+ -- | Subtree's monoid product+ m a+prodSubtree lct@Lct {nLct, subtreeProdLct} v rootOrParent = stToPrim $ do+ if v == rootOrParent+ then do+ -- `v` will be the root+ evertST lct v+ VGM.unsafeRead subtreeProdLct v+ else do+ -- @rootOrParent@ can be far. retrieve the adjacent vertex:+ parent_ <- jumpST lct v rootOrParent 1+ -- detach @v@ from the parent. now that it's the root of the subtree vertices, the aggregation+ -- value is the aggregation of all the subtree vertices.+ cutST lct v parent_+ res <- VGM.unsafeRead subtreeProdLct v+ -- attach again+ linkST lct v parent_+ pure res+ where+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.prodSubtree" v nLct+ !_ = ACIA.checkIndex "AtCoder.Extra.Lct.prodSubtree" rootOrParent nLct
src/AtCoder/Extra/Tree/TreeMonoid.hs view
@@ -163,6 +163,7 @@ ) -- | \(O(n)\)+{-# INLINE buildImpl #-} buildImpl :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Hld.Hld ->@@ -181,6 +182,7 @@ -- | \(O(n)\) Creates a `TreeMonoid` with weights on vertices. -- -- @since 1.1.0.0+{-# INLINE fromVerts #-} fromVerts :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => -- | `Hld.Hld`.@@ -204,6 +206,7 @@ -- duplicated: only one of \((u, v, w)\) or \((v, u, w)\) is needed. -- -- @since 1.1.0.0+{-# INLINE fromEdges #-} fromEdges :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => -- | `Hld.Hld`.@@ -227,6 +230,7 @@ -- | \(O(\log^2 n)\) Returns the product of the path between two vertices \(u\), \(v\) (invlusive). -- -- @since 1.1.0.0+{-# INLINE prod #-} prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> Vertex -> m a prod TreeMonoid {..} u v = do case commuteTM of@@ -236,6 +240,7 @@ -- | \(O(\log n)\) Returns the product of the subtree rooted at the given `Vertex`. -- -- @since 1.1.0.0+{-# INLINE prodSubtree #-} prodSubtree :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a prodSubtree TreeMonoid {..} subtreeRoot = do let (!l, !r) = Hld.subtreeSegmentInclusive hldTM subtreeRoot@@ -247,17 +252,19 @@ then pure mempty else ST.prod segFTM (l + 1) (r + 1) --- | \(O(1)\) Reads a `TreeMonoid` value on a `Vertex`.+-- | \(O(1)\) Reads a monoid value of a `Vertex`. -- -- @since 1.1.0.0+{-# INLINE read #-} read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a read TreeMonoid {..} i_ = do let !i = Hld.indexHld hldTM VG.! i_ ST.read segFTM i --- | \(O(\log n)\) Write a `TreeMonoid` value on a `Vertex`.+-- | \(O(\log n)\) Writes to the monoid value of a vertex. -- -- @since 1.1.0.0+{-# INLINE write #-} write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m () write TreeMonoid {..} i_ x = do let !i = Hld.indexHld hldTM VG.! i_@@ -265,9 +272,10 @@ when (commuteTM == NonCommute) $ do ST.write segBTM i $ Dual x --- | \(O(\log n)\) Exchanges a `TreeMonoid` value on a `Vertex`.+-- | \(O(\log n)\) Writes to the monoid value of a vertex and returns the old value. -- -- @since 1.1.0.0+{-# INLINE exchange #-} exchange :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m a exchange TreeMonoid {..} i_ x = do let !i = Hld.indexHld hldTM VG.! i_@@ -276,9 +284,10 @@ ST.write segBTM i $ Dual x pure res --- | \(O(\log n)\) Modifies a `TreeMonoid` value on a `Vertex`.+-- | \(O(\log n)\) Modifies the monoid value of a vertex with a pure function. -- -- @since 1.1.0.0+{-# INLINE modify #-} modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> (a -> a) -> Int -> m () modify TreeMonoid {..} f i_ = do let !i = Hld.indexHld hldTM VG.! i_@@ -286,9 +295,10 @@ when (commuteTM == NonCommute) $ do ST.modify segBTM (Dual . f . getDual) i --- | \(O(\log n)\) Modifies a `TreeMonoid` value on a `Vertex`.+-- | \(O(\log n)\) Modifies the monoid value of a vertex with a monadic function. -- -- @since 1.1.0.0+{-# INLINE modifyM #-} modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> (a -> m a) -> Int -> m () modifyM TreeMonoid {..} f i_ = do let !i = Hld.indexHld hldTM VG.! i_
src/AtCoder/Extra/WaveletMatrix.hs view
@@ -99,14 +99,14 @@ -- original array if you can. -- -- @since 1.1.0.0-{-# INLINE access #-}+{-# INLINABLE access #-} access :: WaveletMatrix -> Int -> Maybe Int access WaveletMatrix {..} i = (xDictWM VG.!) <$> Rwm.access rawWM i -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\). -- -- @since 1.1.0.0-{-# INLINE rank #-}+{-# INLINABLE rank #-} rank :: -- | A wavelet matrix WaveletMatrix ->@@ -123,7 +123,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\). -- -- @since 1.1.0.0-{-# INLINE rankBetween #-}+{-# INLINABLE rankBetween #-} rankBetween :: -- | A wavelet matrix WaveletMatrix ->@@ -151,7 +151,7 @@ -- not found. -- -- @since 1.1.0.0-{-# INLINE select #-}+{-# INLINABLE select #-} select :: WaveletMatrix -> Int -> Maybe Int select wm = selectKth wm 0 @@ -159,7 +159,7 @@ -- if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINE selectKth #-}+{-# INLINABLE selectKth #-} selectKth :: -- | A wavelet matrix WaveletMatrix ->@@ -180,7 +180,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINE selectIn #-}+{-# INLINABLE selectIn #-} selectIn :: -- | A wavelet matrix WaveletMatrix ->@@ -198,7 +198,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINE selectKthIn #-}+{-# INLINABLE selectKthIn #-} selectKthIn :: -- | A wavelet matrix WaveletMatrix ->@@ -223,7 +223,7 @@ -- largest value. Note that duplicated values are treated as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINE kthLargestIn #-}+{-# INLINABLE kthLargestIn #-} kthLargestIn :: -- | A wavelet matrix WaveletMatrix ->@@ -243,7 +243,7 @@ -- \(k\)-th (0-based) largest value. Note that duplicated values are treated as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINE ikthLargestIn #-}+{-# INLINABLE ikthLargestIn #-} ikthLargestIn :: -- | A wavelet matrix WaveletMatrix ->@@ -263,7 +263,7 @@ -- smallest value. Note that duplicated values are treated as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINE kthSmallestIn #-}+{-# INLINABLE kthSmallestIn #-} kthSmallestIn :: -- | A wavelet matrix WaveletMatrix ->@@ -283,7 +283,7 @@ -- \(k\)-th (0-based) smallest value. Note that duplicated values are treated as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINE ikthSmallestIn #-}+{-# INLINABLE ikthSmallestIn #-} ikthSmallestIn :: WaveletMatrix -> -- | \(l\)@@ -301,7 +301,7 @@ -- | \(O(\log |S|)\) -- -- @since 1.1.0.0-{-# INLINE unsafeKthSmallestIn #-}+{-# INLINABLE unsafeKthSmallestIn #-} unsafeKthSmallestIn :: WaveletMatrix -> Int -> Int -> Int -> Int unsafeKthSmallestIn WaveletMatrix {..} l r k = xDictWM VG.! Rwm.unsafeKthSmallestIn rawWM l r k@@ -309,7 +309,7 @@ -- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0]\). -- -- @since 1.1.0.0-{-# INLINE lookupLE #-}+{-# INLINABLE lookupLE #-} lookupLE :: -- | A wavelet matrix WaveletMatrix ->@@ -334,7 +334,7 @@ -- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0)\). -- -- @since 1.1.0.0-{-# INLINE lookupLT #-}+{-# INLINABLE lookupLT #-} lookupLT :: -- | A wavelet matrix WaveletMatrix ->@@ -351,7 +351,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times [y_0, \infty)\). -- -- @since 1.1.0.0-{-# INLINE lookupGE #-}+{-# INLINABLE lookupGE #-} lookupGE :: -- | A wavelet matrix WaveletMatrix ->@@ -376,7 +376,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\). -- -- @since 1.1.0.0-{-# INLINE lookupGT #-}+{-# INLINABLE lookupGT #-} lookupGT :: -- | A wavelet matrix WaveletMatrix ->@@ -394,7 +394,7 @@ -- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINE assocsIn #-}+{-# INLINABLE assocsIn #-} assocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)] assocsIn WaveletMatrix {..} l r = Rwm.assocsWith rawWM l r (xDictWM VG.!) @@ -402,6 +402,6 @@ -- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINE descAssocsIn #-}+{-# INLINABLE descAssocsIn #-} descAssocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)] descAssocsIn WaveletMatrix {..} l r = Rwm.descAssocsInWith rawWM l r (xDictWM VG.!)
src/AtCoder/Extra/WaveletMatrix/Raw.hs view
@@ -157,7 +157,7 @@ -- original array if you can. -- -- @since 1.1.0.0-{-# INLINE access #-}+{-# INLINABLE access #-} access :: RawWaveletMatrix -> Int -> Maybe Int access RawWaveletMatrix {..} i0 | ACIA.testIndex i0 lengthRwm =@@ -181,7 +181,7 @@ -- | \(O(\log |A|)\) Goes down the wavelet matrix for collecting the kth smallest value. -- -- @since 1.1.0.0-{-# INLINE goDown #-}+{-# INLINABLE goDown #-} goDown :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int, Int, Int) goDown RawWaveletMatrix {..} l_ r_ k_ = V.ifoldl' step (0 :: Int, l_, r_, k_) bitsRwm where@@ -207,7 +207,7 @@ -- | \(O(\log |A|)\) Goes up the wavelet matrix for collecting the value \(x\). -- -- @since 1.1.0.0-{-# INLINE goUp #-}+{-# INLINABLE goUp #-} goUp :: RawWaveletMatrix -> Int -> Int -> Maybe Int goUp RawWaveletMatrix {..} i0 x = V.ifoldM'@@ -222,7 +222,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [0, y_0)\). -- -- @since 1.1.0.0-{-# INLINE rankLT #-}+{-# INLINABLE rankLT #-} rankLT :: RawWaveletMatrix -> Int -> Int -> Int -> Int rankLT RawWaveletMatrix {..} l_ r_ xr -- REMARK: This is required. The function below cannot handle the case N = 2^i and xr = N.@@ -248,7 +248,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\). -- -- @since 1.1.0.0-{-# INLINE rank #-}+{-# INLINABLE rank #-} rank :: RawWaveletMatrix -> -- | \(l\)@@ -264,7 +264,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\). -- -- @since 1.1.0.0-{-# INLINE rankBetween #-}+{-# INLINABLE rankBetween #-} rankBetween :: RawWaveletMatrix -> -- | \(l\)@@ -283,7 +283,7 @@ -- not found. -- -- @since 1.1.0.0-{-# INLINE select #-}+{-# INLINABLE select #-} select :: RawWaveletMatrix -> Int -> Maybe Int select wm = selectKth wm 0 @@ -291,7 +291,7 @@ -- if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINE selectKth #-}+{-# INLINABLE selectKth #-} selectKth :: RawWaveletMatrix -> -- | \(k\)@@ -306,7 +306,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINE selectIn #-}+{-# INLINABLE selectIn #-} selectIn :: -- | A wavelet matrix RawWaveletMatrix ->@@ -324,7 +324,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINE selectKthIn #-}+{-# INLINABLE selectKthIn #-} selectKthIn :: RawWaveletMatrix -> -- | \(l\)@@ -369,7 +369,7 @@ -- largest value. Note that duplicated values are counted as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINE kthLargestIn #-}+{-# INLINABLE kthLargestIn #-} kthLargestIn :: -- | A wavelet matrix RawWaveletMatrix ->@@ -390,7 +390,7 @@ -- \(k\)-th (0-based) largest value. Note that duplicated values are counted as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINE ikthLargestIn #-}+{-# INLINABLE ikthLargestIn #-} ikthLargestIn :: -- | A wavelet matrix RawWaveletMatrix ->@@ -411,7 +411,7 @@ -- smallest value. Note that duplicated values are counted as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINE kthSmallestIn #-}+{-# INLINABLE kthSmallestIn #-} kthSmallestIn :: -- | A wavelet matrix RawWaveletMatrix ->@@ -432,7 +432,7 @@ -- \(k\)-th (0-based) smallest value. Note that duplicated values are counted as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINE ikthSmallestIn #-}+{-# INLINABLE ikthSmallestIn #-} ikthSmallestIn :: RawWaveletMatrix -> -- | \(l\)@@ -452,21 +452,21 @@ -- values are counted as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINE unsafeKthLargestIn #-}+{-# INLINABLE unsafeKthLargestIn #-} unsafeKthLargestIn :: RawWaveletMatrix -> Int -> Int -> Int -> Int unsafeKthLargestIn wm l r k = unsafeKthSmallestIn wm l r (r - l - (k + 1)) -- | \(O(\log a)\) -- -- @since 1.1.0.0-{-# INLINE unsafeIKthLargestIn #-}+{-# INLINABLE unsafeIKthLargestIn #-} unsafeIKthLargestIn :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int) unsafeIKthLargestIn wm l r k = unsafeIKthSmallestIn wm l r (r - l - (k + 1)) -- | \(O(\log a)\) -- -- @since 1.1.0.0-{-# INLINE unsafeKthSmallestIn #-}+{-# INLINABLE unsafeKthSmallestIn #-} unsafeKthSmallestIn :: RawWaveletMatrix -> Int -> Int -> Int -> Int unsafeKthSmallestIn wm l_ r_ k_ = let (!x, !_, !_, !_) = goDown wm l_ r_ k_@@ -475,7 +475,7 @@ -- | \(O(\log a)\) -- -- @since 1.1.0.0-{-# INLINE unsafeIKthSmallestIn #-}+{-# INLINABLE unsafeIKthSmallestIn #-} unsafeIKthSmallestIn :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int) unsafeIKthSmallestIn wm l_ r_ k_ = let (!x, !l, !_, !k) = goDown wm l_ r_ k_@@ -485,7 +485,7 @@ -- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0]\). -- -- @since 1.1.0.0-{-# INLINE lookupLE #-}+{-# INLINABLE lookupLE #-} lookupLE :: -- | A wavelet matrix RawWaveletMatrix ->@@ -510,7 +510,7 @@ -- | \(O(\log a)\) Finds the maximum \(x\) in \([l, r)\) s.t. \(x_{0} \lt x\). -- -- @since 1.1.0.0-{-# INLINE lookupLT #-}+{-# INLINABLE lookupLT #-} lookupLT :: RawWaveletMatrix -> -- | \(l\)@@ -526,7 +526,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times [y_0, \infty)\). -- -- @since 1.1.0.0-{-# INLINE lookupGE #-}+{-# INLINABLE lookupGE #-} lookupGE :: RawWaveletMatrix -> -- | \(l\)@@ -551,7 +551,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\). -- -- @since 1.1.0.0-{-# INLINE lookupGT #-}+{-# INLINABLE lookupGT #-} lookupGT :: RawWaveletMatrix -> -- | \(l\)@@ -568,14 +568,14 @@ -- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINE assocsIn #-}+{-# INLINABLE assocsIn #-} assocsIn :: RawWaveletMatrix -> Int -> Int -> [(Int, Int)] assocsIn wm l r = assocsWith wm l r id -- | \(O(\log A \min(|A|, L))\) Internal implementation of `assocs`. -- -- @since 1.1.0.0-{-# INLINE assocsWith #-}+{-# INLINABLE assocsWith #-} assocsWith :: RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)] assocsWith RawWaveletMatrix {..} l_ r_ f | l'_ < r'_ = inner (0 :: Int) (0 :: Int) l'_ r'_ []@@ -611,14 +611,14 @@ -- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINE descAssocsIn #-}+{-# INLINABLE descAssocsIn #-} descAssocsIn :: RawWaveletMatrix -> Int -> Int -> [(Int, Int)] descAssocsIn wm l r = descAssocsInWith wm l r id -- | \(O(\log A \min(|A|, L))\) Internal implementation of `descAssoc`. -- -- @since 1.1.0.0-{-# INLINE descAssocsInWith #-}+{-# INLINABLE descAssocsInWith #-} descAssocsInWith :: RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)] descAssocsInWith RawWaveletMatrix {..} l_ r_ f | l'_ < r'_ = inner (0 :: Int) (0 :: Int) l'_ r'_ []
src/AtCoder/Internal/Buffer.hs view
@@ -65,6 +65,7 @@ -- * Modifications pushBack, popBack,+ popBack_, write, modify, modifyM,@@ -181,6 +182,15 @@ x <- VGM.read vecB (len - 1) VGM.write lenB 0 (len - 1) pure $ Just x++-- | \(O(1)\) Removes the last element from the buffer and discards it.+--+-- @since 1.1.1.0+{-# INLINE popBack_ #-}+popBack_ :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m ()+popBack_ buf = do+ _ <- popBack buf+ pure () -- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is -- out of bounds.
src/AtCoder/Internal/Convolution.hs view
@@ -253,13 +253,13 @@ ans <- VGM.replicate (n + m - 1) 0 if n < m then do- for_ [0 .. m - 1] $ \j -> do- for_ [0 .. n - 1] $ \i -> do- VGM.modify ans (+ a VG.! i * b VG.! j) (i + j)+ VU.iforM_ b $ \j bj -> do+ VU.iforM_ a $ \i ai -> do+ VGM.modify ans (+ ai * bj) (i + j) else do- for_ [0 .. n - 1] $ \i -> do- for_ [0 .. m - 1] $ \j -> do- VGM.modify ans (+ a VG.! i * b VG.! j) (i + j)+ VU.iforM_ a $ \i ai -> do+ VU.iforM_ b $ \j bj -> do+ VGM.modify ans (+ ai * bj) (i + j) pure ans -- | @since 1.0.0.0
src/AtCoder/Internal/Math.hs view
@@ -52,7 +52,7 @@ -- -- ==== Constraints -- - \(0 \le n\)--- - \(1 \le m\)+-- - \(1 \le m \lt 2^{31}\) -- -- ==== Complexity -- - \(O(\log n)\)@@ -145,6 +145,10 @@ in inner t s' m1 m0' -- | Returns the primitive root of the given `Int`.+--+-- ==== Constraints+-- - The input must be a prime number.+-- - The input must be less than \(2^31\). -- -- @since 1.0.0.0 {-# INLINE primitiveRoot #-}
src/AtCoder/Internal/Scc.hs view
@@ -98,7 +98,7 @@ -- | \(O(n + m)\) API) Returns a pair of @(# of scc, scc id)@. -- -- @since 1.1.0.0-{-# INLINE sccIdsCsr #-}+{-# INLINABLE sccIdsCsr #-} sccIdsCsr :: ACICSR.Csr w -> (Int, VU.Vector Int) sccIdsCsr g@ACICSR.Csr {..} = runST $ do -- see also the Wikipedia: https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm#The_algorithm_in_pseudocode@@ -173,7 +173,7 @@ -- | \(O(n + m)\) Returns the strongly connected components. -- -- @since 1.1.0.0-{-# INLINE sccCsr #-}+{-# INLINABLE sccCsr #-} sccCsr :: ACICSR.Csr w -> V.Vector (VU.Vector Int) sccCsr g = runST $ do groups <- V.mapM VUM.unsafeNew $ VU.convert counts
src/AtCoder/Internal/String.hs view
@@ -25,7 +25,7 @@ -- | \(O(n^2)\) Internal implementation of suffix array creation (naive). -- -- @since 1.0.0.0-{-# INLINE saNaive #-}+{-# INLINABLE saNaive #-} saNaive :: (HasCallStack) => VU.Vector Int -> VU.Vector Int saNaive s = let n = VU.length s@@ -48,7 +48,7 @@ -- | \(O(n \log n)\) Internal implementation of suffix array creation (doubling). -- -- @since 1.0.0.0-{-# INLINE saDoubling #-}+{-# INLINABLE saDoubling #-} saDoubling :: (HasCallStack) => VU.Vector Int -> VU.Vector Int saDoubling s = VU.create $ do let n = VU.length s@@ -87,7 +87,7 @@ -- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting). -- -- @since 1.0.0.0-{-# INLINE saIsImpl #-}+{-# INLINABLE saIsImpl #-} saIsImpl :: (HasCallStack) => -- | naive threshould
src/AtCoder/String.hs view
@@ -101,7 +101,7 @@ -- - \(O(n)\)-space -- -- @since 1.0.0.0-{-# INLINE suffixArrayOrd #-}+{-# INLINABLE suffixArrayOrd #-} suffixArrayOrd :: (HasCallStack, Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int suffixArrayOrd s = let n = VU.length s@@ -137,7 +137,7 @@ -- - \(O(n)\) -- -- @since 1.0.0.0-{-# INLINE lcpArray #-}+{-# INLINABLE lcpArray #-} lcpArray :: (HasCallStack, Ord a, VU.Unbox a) => -- | A vector representing a string@@ -209,7 +209,7 @@ -- - \(O(n)\) -- -- @since 1.0.0.0-{-# INLINE zAlgorithm #-}+{-# INLINABLE zAlgorithm #-} zAlgorithm :: (Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int zAlgorithm s | n == 0 = VU.empty
test/Tests/Extra/HashMap.hs view
@@ -59,23 +59,22 @@ instance QC.Arbitrary Query where arbitrary = do- a <- QC.chooseInt (1, 100)- if a == 1- then pure Clear- else- QC.oneof- [ pure Size,- Member <$> keyGen,- NotMember <$> keyGen,- Lookup <$> keyGen,- Insert <$> keyGen <*> valGen,- InsertWithAdd <$> keyGen <*> valGen,- Exchange <$> keyGen <*> valGen,- ModifyAdd <$> keyGen <*> valGen- ]+ QC.frequency+ [ (rare, pure Clear),+ (often, pure Size),+ (often, Member <$> keyGen),+ (often, NotMember <$> keyGen),+ (often, Lookup <$> keyGen),+ (often, Insert <$> keyGen <*> valGen),+ (often, InsertWithAdd <$> keyGen <*> valGen),+ (often, Exchange <$> keyGen <*> valGen),+ (often, ModifyAdd <$> keyGen <*> valGen)+ ] where+ rare = 1+ often = 10 keyGen = QC.chooseInt (-5, 5)- valGen = QC.chooseInt (-10, 10)+ valGen = QC.arbitrary @Int -- | Arbitrary return type for the `Query` result. data Result
test/Tests/Extra/Semigroup/Matrix.hs view
@@ -2,6 +2,7 @@ import AtCoder.Extra.Semigroup.Matrix qualified as Mat import AtCoder.ModInt qualified as M+import Data.Semigroup (stimes) import Data.Vector.Unboxed qualified as VU import GHC.TypeNats (KnownNat) import Test.QuickCheck.Classes qualified as QCC@@ -34,10 +35,22 @@ let rhs = Mat.vecM $ Mat.mul mat (Mat.new w 1 col) pure $ lhs QC.=== rhs +m :: Int+m = 998244353++prop_pow :: QC.Small Int -> QC.NonNegative Int -> QC.Gen QC.Property+prop_pow (QC.Small n) (QC.NonNegative k) = do+ vec <- VU.fromList <$> QC.vectorOf (n * n) (QC.chooseInt (0, m - 1))+ let mat = Mat.new n n vec+ if k == 0+ then pure $ Mat.pow k mat QC.=== Mat.ident n+ else pure $ Mat.pow k mat QC.=== stimes k mat+ tests :: [TestTree] tests = [ QC.testProperty "mulToCol" prop_mulToCol, laws @(Mat.Matrix (M.ModInt 998244353)) [ QCC.semigroupLaws- ]+ ],+ QC.testProperty "pow" prop_pow ]