diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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`. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -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 diff --git a/ac-library-hs.cabal b/ac-library-hs.cabal --- a/ac-library-hs.cabal +++ b/ac-library-hs.cabal @@ -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 diff --git a/benchmarks/Bench/Matrix.hs b/benchmarks/Bench/Matrix.hs --- a/benchmarks/Bench/Matrix.hs +++ b/benchmarks/Bench/Matrix.hs @@ -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) diff --git a/benchmarks/BenchLib/Matrix.hs b/benchmarks/BenchLib/Matrix.hs --- a/benchmarks/BenchLib/Matrix.hs +++ b/benchmarks/BenchLib/Matrix.hs @@ -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" diff --git a/benchmarks/Main.hs b/benchmarks/Main.hs --- a/benchmarks/Main.hs +++ b/benchmarks/Main.hs @@ -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 diff --git a/src/AtCoder/Convolution.hs b/src/AtCoder/Convolution.hs --- a/src/AtCoder/Convolution.hs +++ b/src/AtCoder/Convolution.hs @@ -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 -> diff --git a/src/AtCoder/Dsu.hs b/src/AtCoder/Dsu.hs --- a/src/AtCoder/Dsu.hs +++ b/src/AtCoder/Dsu.hs @@ -46,7 +46,7 @@ ( -- * Disjoint set union Dsu (nDsu), - -- * Constructor + -- * Constructors new, -- * Merging diff --git a/src/AtCoder/Extra/Graph.hs b/src/AtCoder/Extra/Graph.hs --- a/src/AtCoder/Extra/Graph.hs +++ b/src/AtCoder/Extra/Graph.hs @@ -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)) diff --git a/src/AtCoder/Extra/IntervalMap.hs b/src/AtCoder/Extra/IntervalMap.hs --- a/src/AtCoder/Extra/IntervalMap.hs +++ b/src/AtCoder/Extra/IntervalMap.hs @@ -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 diff --git a/src/AtCoder/Extra/MultiSet.hs b/src/AtCoder/Extra/MultiSet.hs --- a/src/AtCoder/Extra/MultiSet.hs +++ b/src/AtCoder/Extra/MultiSet.hs @@ -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 diff --git a/src/AtCoder/Extra/Pdsu.hs b/src/AtCoder/Extra/Pdsu.hs --- a/src/AtCoder/Extra/Pdsu.hs +++ b/src/AtCoder/Extra/Pdsu.hs @@ -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 diff --git a/src/AtCoder/Extra/Semigroup/Matrix.hs b/src/AtCoder/Extra/Semigroup/Matrix.hs --- a/src/AtCoder/Extra/Semigroup/Matrix.hs +++ b/src/AtCoder/Extra/Semigroup/Matrix.hs @@ -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 diff --git a/src/AtCoder/Extra/Tree/Hld.hs b/src/AtCoder/Extra/Tree/Hld.hs --- a/src/AtCoder/Extra/Tree/Hld.hs +++ b/src/AtCoder/Extra/Tree/Hld.hs @@ -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 diff --git a/src/AtCoder/Extra/Tree/Lct.hs b/src/AtCoder/Extra/Tree/Lct.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Tree/Lct.hs @@ -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 diff --git a/src/AtCoder/Extra/Tree/TreeMonoid.hs b/src/AtCoder/Extra/Tree/TreeMonoid.hs --- a/src/AtCoder/Extra/Tree/TreeMonoid.hs +++ b/src/AtCoder/Extra/Tree/TreeMonoid.hs @@ -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_ diff --git a/src/AtCoder/Extra/WaveletMatrix.hs b/src/AtCoder/Extra/WaveletMatrix.hs --- a/src/AtCoder/Extra/WaveletMatrix.hs +++ b/src/AtCoder/Extra/WaveletMatrix.hs @@ -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.!) diff --git a/src/AtCoder/Extra/WaveletMatrix/Raw.hs b/src/AtCoder/Extra/WaveletMatrix/Raw.hs --- a/src/AtCoder/Extra/WaveletMatrix/Raw.hs +++ b/src/AtCoder/Extra/WaveletMatrix/Raw.hs @@ -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'_ [] diff --git a/src/AtCoder/Internal/Buffer.hs b/src/AtCoder/Internal/Buffer.hs --- a/src/AtCoder/Internal/Buffer.hs +++ b/src/AtCoder/Internal/Buffer.hs @@ -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. diff --git a/src/AtCoder/Internal/Convolution.hs b/src/AtCoder/Internal/Convolution.hs --- a/src/AtCoder/Internal/Convolution.hs +++ b/src/AtCoder/Internal/Convolution.hs @@ -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 diff --git a/src/AtCoder/Internal/Math.hs b/src/AtCoder/Internal/Math.hs --- a/src/AtCoder/Internal/Math.hs +++ b/src/AtCoder/Internal/Math.hs @@ -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 #-} diff --git a/src/AtCoder/Internal/Scc.hs b/src/AtCoder/Internal/Scc.hs --- a/src/AtCoder/Internal/Scc.hs +++ b/src/AtCoder/Internal/Scc.hs @@ -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 diff --git a/src/AtCoder/Internal/String.hs b/src/AtCoder/Internal/String.hs --- a/src/AtCoder/Internal/String.hs +++ b/src/AtCoder/Internal/String.hs @@ -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 diff --git a/src/AtCoder/String.hs b/src/AtCoder/String.hs --- a/src/AtCoder/String.hs +++ b/src/AtCoder/String.hs @@ -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 diff --git a/test/Tests/Extra/HashMap.hs b/test/Tests/Extra/HashMap.hs --- a/test/Tests/Extra/HashMap.hs +++ b/test/Tests/Extra/HashMap.hs @@ -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 diff --git a/test/Tests/Extra/Semigroup/Matrix.hs b/test/Tests/Extra/Semigroup/Matrix.hs --- a/test/Tests/Extra/Semigroup/Matrix.hs +++ b/test/Tests/Extra/Semigroup/Matrix.hs @@ -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 ]