diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,16 @@ # Revision history for acl-hs +## 1.1.0.0 -- Jan 2025 + +- Removed `RangeSetId` and `RangeAddId` from `AtCoder.Extra.Monoid`. +- Implemented `SegAct` for `RangeSet`, `RangeAdd` and `Max`, `Min`. +- Added `segActWithLength` to `SegAct`. +- Added `build1` to `AtCoder.Internal.Csr`. +- Added a bunch of extra modules. + ## 1.0.0.0 -- Dec 2024 -* First version. -* ACL-compatible modules. -* Extra module of `Math` (binary exponentiation) and `Monoid` (`SegAct` instances). +- First version. +- Added ACL-compatible modules. +- Added Extra module of `Math` (binary exponentiation) and `Monoid` (`SegAct` instances). diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -5,8 +5,8 @@ ## Notes - The library is mainly for AtCoder and only GHC 9.8.4 is guaranteed to be supported. -- The `Internal` module is unstable and does not follow PVP. -- The `Extra` module provides with extra contents not in the original `ac-library`. Note that they're less tested. +- 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,12 +4,15 @@ -- PVP summary: +-+------- breaking API changes -- | | +----- non-breaking API additions -- | | | +--- code changes with no API change -version: 1.0.0.1 +version: 1.1.0.0 synopsis: Data structures and algorithms description: Haskell port of [ac-library](https://github.com/atcoder/ac-library), a library for competitive programming on [AtCoder](https://atcoder.jp/). + - Functions primarily use half-open interval \([l, r)\). + - The extra module contains additional utilities beyond the original C++ library. + category: Algorithms, Data Structures license: CC0-1.0 license-file: LICENSE @@ -70,13 +73,31 @@ exposed-modules: AtCoder.Convolution AtCoder.Dsu + AtCoder.Extra.Bisect + 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 + AtCoder.Extra.Monoid.Mat2x2 AtCoder.Extra.Monoid.RangeAdd - AtCoder.Extra.Monoid.RangeAddId AtCoder.Extra.Monoid.RangeSet - AtCoder.Extra.Monoid.RangeSetId + AtCoder.Extra.Monoid.RollingHash + AtCoder.Extra.Monoid.V2 + AtCoder.Extra.MultiSet + AtCoder.Extra.Semigroup.Matrix + AtCoder.Extra.Semigroup.Permutation + AtCoder.Extra.Pdsu + AtCoder.Extra.WaveletMatrix + AtCoder.Extra.WaveletMatrix.BitVector + AtCoder.Extra.WaveletMatrix.Raw + AtCoder.Extra.WaveletMatrix2d AtCoder.FenwickTree AtCoder.Internal.Assert AtCoder.Internal.Barrett @@ -118,8 +139,21 @@ other-modules: Tests.Convolution Tests.Dsu + Tests.Extra.Bisect + Tests.Extra.Graph + Tests.Extra.HashMap + Tests.Extra.IntervalMap + Tests.Extra.IntMap + Tests.Extra.IntSet Tests.Extra.Math Tests.Extra.Monoid + Tests.Extra.MultiSet + Tests.Extra.Semigroup.Matrix + Tests.Extra.Semigroup.Permutation + Tests.Extra.WaveletMatrix + Tests.Extra.WaveletMatrix.BitVector + Tests.Extra.WaveletMatrix.Raw + Tests.Extra.WaveletMatrix2d Tests.FenwickTree Tests.Internal.Bit Tests.Internal.Buffer @@ -146,6 +180,7 @@ main-is: Main.hs build-depends: , ac-library-hs + , containers , hspec , mtl , QuickCheck @@ -157,6 +192,7 @@ , tasty-quickcheck , tasty-rerun , transformers + , unordered-containers benchmark ac-library-hs-benchmark import: warnings @@ -167,8 +203,11 @@ hs-source-dirs: benchmarks other-modules: Bench.AddMod + Bench.Matrix + Bench.ModInt Bench.PowMod BenchLib.AddMod + BenchLib.Matrix BenchLib.ModInt.ModIntNats BenchLib.ModInt.Modulus BenchLib.MulMod.Barrett64 @@ -180,8 +219,10 @@ , ac-library-hs , base , criterion + , mtl , random , tagged + , transformers test-suite benchlib-test import: warnings @@ -204,6 +245,7 @@ , hspec , mtl , QuickCheck + -- , quickcheck-instances , random , tagged , tasty diff --git a/benchmarks/Bench/AddMod.hs b/benchmarks/Bench/AddMod.hs --- a/benchmarks/Bench/AddMod.hs +++ b/benchmarks/Bench/AddMod.hs @@ -6,7 +6,7 @@ import BenchLib.AddMod import Criterion import Data.Vector.Unboxed qualified as VU -import Data.Word (Word32, Word64) +import Data.Word (Word32) import System.Random n :: Int diff --git a/benchmarks/Bench/Matrix.hs b/benchmarks/Bench/Matrix.hs new file mode 100644 --- /dev/null +++ b/benchmarks/Bench/Matrix.hs @@ -0,0 +1,87 @@ +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 BenchLib.Matrix qualified as Mat +import Control.Monad.State.Class (MonadState, state) +import Control.Monad.Trans.State.Strict (evalState, runState) +import Criterion +import Data.Semigroup (Semigroup (..)) +import Data.Vector qualified as V +import Data.Vector.Unboxed qualified as VU +import System.Random + +testSize :: Int +testSize = 10000 + +m :: Int +m = 998244353 + +randomMatrix :: (MonadState StdGen m) => m (Mat.Matrix Int, Int) +randomMatrix = do + h <- state $ uniformR (1, 16) + nextMatrix h + +nextMatrix :: (MonadState StdGen m) => Int -> m (Mat.Matrix Int, Int) +nextMatrix h = do + w <- state $ uniformR (1, 16) + vec <- VU.replicateM (h * w) $ state (uniformR (0, m - 1)) + pure (Mat.new h w vec, w) + +randomSquareACLMatrix :: (MonadState StdGen m) => Int -> m (ACMAT.Matrix Int) +randomSquareACLMatrix n = do + vec <- VU.replicateM (n * n) $ state (uniformR (0, m - 1)) + pure $ ACMAT.new n n vec + +benches :: Benchmark +benches = + bgroup + "Matrix" + [ bench "mul1" $ whnf (V.foldl1' Mat.mul1) randomMatrixInput, + bench "mul2" $ whnf (V.foldl1' Mat.mul2) randomMatrixInput, + 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 "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 "mulMint1" $ whnf (V.foldl1' Mat.mulMint1) randomMintMatrixInput, + bench "mulMint2" $ whnf (V.foldl1' Mat.mulMint2) randomMintMatrixInput, + bench "mulMint3" $ whnf (V.foldl1' Mat.mulMint3) randomMintMatrixInput, + -- mul (ACL) + bench "mul_ACL" $ whnf (V.foldl1' ACMAT.mul) randomMatrixInputACL, + bench "mulMod_ACL" $ whnf (V.foldl1' (ACMAT.mulMod m)) randomMatrixInputACL, + -- 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, + bench "powMintACL_stimes'" $ whnf (VU.foldl' (flip ACEM.stimes') squareMatMint) randomVec, + bench "powMintACL_powMint" $ whnf (VU.foldl' (flip ACMAT.powMint) squareMatMint) randomVec + ] + where + -- Bench matrix + randomMatrixInput :: V.Vector (Mat.Matrix Int) + !randomMatrixInput = evalState (V.unfoldrExactNM testSize nextMatrix (Mat.wM mat0)) gen0 + where + ((!mat0, !_), !gen0) = runState randomMatrix $ mkStdGen 123456789 + + randomMintMatrixInput :: V.Vector (Mat.Matrix (M.ModInt 998244353)) + !randomMintMatrixInput = V.map (Mat.map M.new) randomMatrixInput + + -- ACL matrix + randomMatrixInputACL :: V.Vector (ACMAT.Matrix Int) + !randomMatrixInputACL = V.map (\mat -> ACMAT.new (Mat.hM mat) (Mat.wM mat) (Mat.vecM mat)) randomMatrixInput + + squareMat :: ACMAT.Matrix Int + !squareMat = evalState (randomSquareACLMatrix 17) $ mkStdGen 123456789 + + 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) diff --git a/benchmarks/Bench/ModInt.hs b/benchmarks/Bench/ModInt.hs new file mode 100644 --- /dev/null +++ b/benchmarks/Bench/ModInt.hs @@ -0,0 +1,23 @@ +module Bench.ModInt (benches) where + +import AtCoder.ModInt qualified as M +import BenchLib.PowMod qualified as PowMod +import BenchLib.ModInt.ModIntNats qualified as MN +import Criterion +import Data.Vector.Unboxed qualified as VU +import System.Random + +benches :: Benchmark +benches = + bgroup + "modInt" + [ + + ] + where + n = 10000 + randomVec :: VU.Vector Int + randomVec = + VU.map fromIntegral $ + VU.unfoldrExactN n (genWord64R (998244383 - 2)) (mkStdGen 123456789) + diff --git a/benchmarks/BenchLib/Matrix.hs b/benchmarks/BenchLib/Matrix.hs new file mode 100644 --- /dev/null +++ b/benchmarks/BenchLib/Matrix.hs @@ -0,0 +1,394 @@ +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RecordWildCards #-} + +module BenchLib.Matrix + ( Matrix (..), + new, + map, + mulToCol, + mulToColModInt, + mul1, + mul2, + mul3_1, + mul3_2, + mul3_3, + mulMod1, + mulMod2, + mulMod3, + mulMod4, + mulMod5, + mulMint1, + mulMint2, + mulMint3, + ) +where + +import AtCoder.Internal.Assert qualified as ACIA +import AtCoder.Internal.Barrett qualified as BT +import AtCoder.ModInt qualified as M +import Data.Vector qualified as V +import Data.Vector.Generic qualified as VG +import Data.Vector.Unboxed qualified as VU +import Data.Word (Word64) +import GHC.Exts (Proxy#, proxy#) +import GHC.Stack (HasCallStack) +import GHC.TypeNats (KnownNat, natVal, natVal') +import Prelude hiding (map) + +data Matrix a = Matrix + { hM :: {-# UNPACK #-} !Int, + wM :: {-# UNPACK #-} !Int, + vecM :: !(VU.Vector a) + } + deriving (Show, Eq) + +type Col a = VU.Vector a + +{-# INLINE new #-} +new :: (HasCallStack, VU.Unbox a) => Int -> Int -> VU.Vector a -> Matrix a +new h w vec + | VU.length vec /= h * w = error "AtCoder.Extra.Matrix: size mismatch" + | otherwise = Matrix h w vec + +{-# INLINE map #-} +map :: (VU.Unbox a, VU.Unbox b) => (a -> b) -> Matrix a -> Matrix b +map f Matrix {..} = Matrix hM wM $ VU.map f vecM + +{-# INLINE mulToCol #-} +mulToCol :: (Num a, VU.Unbox a) => Matrix a -> Col a -> Col a +mulToCol Matrix {..} !col = VU.convert $ V.map (VU.sum . VU.zipWith (*) col) rows + where + !n = VU.length col + !_ = ACIA.runtimeAssert (n == wM) "AtCoder.Extra.Matrix.mulToCol: size mismatch" + rows = V.unfoldrExactN hM (VU.splitAt wM) vecM + +{-# INLINE mulToColModInt #-} +mulToColModInt :: forall m. (KnownNat m) => Matrix (M.ModInt m) -> Col (M.ModInt m) -> Col (M.ModInt m) +mulToColModInt Matrix {..} !col = VU.convert $ V.map (VU.foldl' (+) (M.unsafeNew 0) . VU.zipWith mulMod col) rows + where + !_ = ACIA.runtimeAssert (VU.length col == wM) "AtCoder.Extra.Matrix.mulToColModInt: size mismatch" + !bt = BT.new32 $ fromIntegral (natVal' (proxy# @m)) + rows = V.unfoldrExactN hM (VU.splitAt wM) vecM + mulMod (M.ModInt x) (M.ModInt y) = M.unsafeNew . fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y) + +{-# INLINE mul1 #-} +mul1 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e +mul1 !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) + 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 = + 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 + 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 = + 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 + 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 = + Matrix h w' $ + VU.unfoldrExactN + (h * w') + ( \(!row, !col) -> + let !y = VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (VU.unsafeSlice (w * row) w vecA) + in if col + 1 >= w' + then (y, (row + 1, 0)) + else (y, (row, col + 1)) + ) + (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 -> + 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" + +{-# INLINE mulMod1 #-} +mulMod1 :: Int -> Matrix Int -> Matrix Int -> Matrix Int +mulMod1 !m !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) + where + 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 = + 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 + f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA) + -- very slow + addMod x y + | 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 = + 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) `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 = + 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) `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 = + 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 mulMint1 #-} +mulMint1 :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a) +mulMint1 !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) + where + 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 = (*) + 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 = + 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 (natVal' (proxy# @a)) + 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.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 = + 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 (natVal' (proxy# @a)) + 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) + 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" diff --git a/benchmarks/BenchLib/MulMod/BarrettWideWord.hs b/benchmarks/BenchLib/MulMod/BarrettWideWord.hs --- a/benchmarks/BenchLib/MulMod/BarrettWideWord.hs +++ b/benchmarks/BenchLib/MulMod/BarrettWideWord.hs @@ -24,10 +24,11 @@ imBarrett :: {-# UNPACK #-} !Word64 } --- | Creates barret reduction for modulus \(m\). +-- | Creates a `Barrett` for modulus value \(m\) of type `Word32`. new32 :: Word32 -> Barrett new32 m = Barrett m $ maxBound @Word64 `div` (fromIntegral m :: Word64) + 1 +-- | Creates a `Barrett` for modulus value \(m\) of type `Word64`. new64 :: Word64 -> Barrett new64 m = Barrett (fromIntegral m) $ maxBound @Word64 `div` m + 1 diff --git a/benchmarks/BenchLib/MulMod/Montgomery.hs b/benchmarks/BenchLib/MulMod/Montgomery.hs --- a/benchmarks/BenchLib/MulMod/Montgomery.hs +++ b/benchmarks/BenchLib/MulMod/Montgomery.hs @@ -17,8 +17,8 @@ import Data.Bits (bit, (!>>.)) import Data.Word (Word32, Word64) --- | Fast modular multiplication by Montgomery multiplication. The modulus value has to be odd --- to be fast. +-- | Fast modular multiplication by Montgomery multiplication. The modulus value must be odd for +-- the speed. data Montgomery = Montgomery { mMontgomery :: {-# UNPACK #-} !Word64, -- | R2 == (2^64) % MOD; @@ -27,7 +27,7 @@ negInvMontgomery :: {-# UNPACK #-} !Word32 } --- | Creates `Montgomery` for modulus @m@. +-- | Creates a `Montgomery` for modulus @m@. new :: Word64 -> Montgomery new m = let !negInv = inner 0 0 1 0 diff --git a/benchmarks/Main.hs b/benchmarks/Main.hs --- a/benchmarks/Main.hs +++ b/benchmarks/Main.hs @@ -3,7 +3,9 @@ module Main where import Bench.AddMod qualified +import Bench.Matrix qualified import Bench.MulMod qualified +import Bench.ModInt qualified import Bench.PowMod qualified import Criterion.Main @@ -13,9 +15,10 @@ main :: IO () main = defaultMain - -- TODO: generate graph by benchmark group? - [ -- Bench.MulMod.benches32, - -- Bench.MulMod.benches64, - -- Bench.AddMod.benches - Bench.PowMod.benches + -- TODO: generate criterion graph by benchmark group? + [ Bench.MulMod.benches, + Bench.ModInt.benches, + Bench.AddMod.benches, + Bench.PowMod.benches, + Bench.Matrix.benches ] diff --git a/src/AtCoder/Convolution.hs b/src/AtCoder/Convolution.hs --- a/src/AtCoder/Convolution.hs +++ b/src/AtCoder/Convolution.hs @@ -14,7 +14,7 @@ -- >>> import Data.Proxy (Proxy) -- >>> import Data.Vector.Unboxed qualified as VU -- --- It's handy to define specific items for interested modulus values: +-- Define specific modulus items: -- -- >>> type Mint = M.ModInt998244353 -- >>> let modInt :: Int -> Mint; modInt = M.new @@ -26,7 +26,7 @@ -- >>> C.convolution a b -- [5,16,34,60,61,52,32] -- --- You can also target any @'Integral' a@ (with some runtime overhead for conversion to modint): +-- You can also target any @'Integral' a@ with `convolutionRaw`: -- -- >>> let a = VU.fromList @Int [1, 2, 3, 4] -- >>> let b = VU.fromList @Int [5, 6, 7, 8] diff --git a/src/AtCoder/Dsu.hs b/src/AtCoder/Dsu.hs --- a/src/AtCoder/Dsu.hs +++ b/src/AtCoder/Dsu.hs @@ -1,17 +1,17 @@ {-# LANGUAGE RecordWildCards #-} --- | Disjoint set union, also known as Union-Find tree. It processes the following queries in +-- | A disjoint set union, also known as a Union-Find tree. It processes the following queries in -- amortized \(O(\alpha(n))\) time. -- --- - Edge addition --- - Deciding whether given two vertices are in the same connected component +-- - Edge addition (`merge`) +-- - Deciding whether given two vertices are in the same connected component (`same`) -- --- Each connected component internally has a representative vertex. When two connected components --- are merged by edge addition, one of the two representatives of these connected components --- becomes the representative of the new connected component. +-- Each connected component internally has a representative vertex (`leader`). When two connected +-- components are merged by edge addition (`merge`), one of the two representatives of these +-- connected components becomes the representative (`leader`) of the new connected component. -- -- ==== __Example__ --- Create DSU: +-- Create a `Dsu` with four vertices: -- -- >>> import AtCoder.Dsu qualified as Dsu -- >>> dsu <- Dsu.new 4 -- 0 1 2 3 @@ -49,15 +49,15 @@ -- * Constructor new, - -- * Merge + -- * Merging merge, merge_, -- * Leader leader, - same, -- * Component information + same, size, groups, ) @@ -73,13 +73,13 @@ import Data.Vector.Unboxed.Mutable qualified as VUM import GHC.Stack (HasCallStack) --- | Disjoint set union. Akso known as Union-Find tree. +-- | A disjoint set union. Akso known as Union-Find tree. -- -- @since 1.0.0.0 data Dsu s = Dsu - { -- | 1.0.0 The number of nodes. + { -- | The number of nodes. -- - -- @since + -- @since 1.0.0.0 nDsu :: {-# UNPACK #-} !Int, -- | For root (leader) nodes it stores their size as a negative number. For child nodes it -- stores their parent node index. @@ -103,8 +103,8 @@ pure Dsu {..} | otherwise = error $ "new: given negative size (`" ++ show nDsu ++ "`)" --- | Adds an edge @(a, b)@. the vertices @a@ and @b@ were in the same connected component, it --- returns the representative of this connected component. Otherwise, it returns the +-- | Adds an edge \((a, b)\). If the vertices \(a\) and \(b\) are in the same connected component, it +-- returns the representative (`leader`) of this connected component. Otherwise, it returns the -- representative of the new connected component. -- -- ==== Constraints @@ -134,7 +134,7 @@ VGM.modify parentOrSizeDsu (+ sizeY) x pure x --- | `merge` with return value discarded. +-- | `merge` with the return value discarded. -- -- ==== Constraints -- - \(0 \leq a < n\) @@ -207,7 +207,7 @@ sizeLa <- VGM.read parentOrSizeDsu la pure (-sizeLa) --- | Divides the graph into connected components and returns the list of them. +-- | Divides the graph into connected components and returns the vector of them. -- -- More precisely, it returns a vector of the "vector of the vertices in a connected component". -- Both of the orders of the connected components and the vertices are undefined. diff --git a/src/AtCoder/Extra/Bisect.hs b/src/AtCoder/Extra/Bisect.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Bisect.hs @@ -0,0 +1,297 @@ +-- | Bisection methods and binary search functions. They partition a half-open interval \([l, r)\) +-- into two and return either the left or the right point of the boundary. +-- +-- @ +-- Y Y Y Y Y N N N N N Y: user predicate holds +-- --------* *---------> X N: user predicate does not hold +-- L R L, R: left, right point of the boundary +-- @ +-- +-- ==== __Example__ +-- Perform index compression: +-- +-- >>> import AtCoder.Extra.Bisect +-- >>> import Data.Maybe (fromJust) +-- >>> import Data.Vector.Algorithms.Intro qualified as VAI +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let xs = VU.fromList ([0, 20, 10, 40, 30] :: [Int]) +-- >>> let dict = VU.uniq $ VU.modify VAI.sort xs +-- >>> VU.map (fromJust . lowerBound dict) xs +-- [0,2,1,4,3] +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Bisect + ( -- * C++-like binary search + lowerBound, + lowerBoundIn, + upperBound, + upperBoundIn, + + -- * Generic bisection method + bisectL, + bisectLM, + bisectR, + bisectRM, + ) +where + +import AtCoder.Internal.Assert qualified as ACIA +import Data.Functor ((<&>)) +import Data.Functor.Identity +import Data.Vector.Generic qualified as VG +import GHC.Stack (HasCallStack) + +-- | \(O(\log n)\) Bisection method implementation. Works on a half-open interfal \([l, r)\) . +-- +-- @since 1.1.0.0 +{-# INLINE bisectLImpl #-} +bisectLImpl :: (HasCallStack, Monad m) => (Int -> m Bool) -> Int -> Int -> m Int +bisectLImpl p l0 = inner (l0 - 1) + where + inner l r + | l + 1 == r = pure l + | otherwise = + p mid >>= \case + True -> inner mid r + False -> inner l mid + where + mid = (l + r) `div` 2 + +-- | \(O(\log n)\) Bisection method implementation. Works on a half-open interfal \([l, r)\) . +-- +-- @since 1.1.0.0 +{-# INLINE bisectRImpl #-} +bisectRImpl :: (HasCallStack, Monad m) => (Int -> m Bool) -> Int -> Int -> m Int +bisectRImpl p l = ((+ 1) <$>) . bisectLImpl p l + +-- | \(O(\log n)\) Returns the index of the first element \(x\) in the vector such that +-- \(x \ge x_0\), or `Nothing` if no such element exists. +-- +-- @ +-- Y Y Y Y Y N N N N N Y: (< x0) +-- --------- *---------> X N: (>= x0) +-- R R: returning point +-- @ +-- +-- ==== __Example__ +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let xs = VU.fromList [1, 1, 2, 2, 4, 4] +-- >>> lowerBound xs 1 +-- Just 0 +-- +-- >>> lowerBound xs 2 +-- Just 2 +-- +-- >>> lowerBound xs 3 +-- Just 4 +-- +-- >>> lowerBound xs 4 +-- Just 4 +-- +-- Out of range values: +-- +-- >>> lowerBound xs 0 +-- Just 0 +-- +-- >>> lowerBound xs 5 +-- Nothing +-- +-- @since 1.1.0.0 +{-# INLINE lowerBound #-} +lowerBound :: (HasCallStack, VG.Vector v a, Ord a) => v a -> a -> Maybe Int +lowerBound vec = lowerBoundIn 0 (VG.length vec) vec + +-- | \(O(\log n)\) Computes the `lowerBound` for a slice of a vector within the interval \([l, r)\). +-- +-- - The user predicate evaluates indices in \([l, r)\). +-- - The interval \([l, r)\) is silently clamped to ensure it remains within the bounds \([0, n)\). +-- +-- ==== __Example__ +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let xs = VU.fromList [10, 10, 20, 20, 40, 40] +-- >>> -- *---*---* +-- >>> lowerBoundIn 2 5 xs 10 +-- Just 2 +-- +-- >>> lowerBoundIn 2 5 xs 20 +-- Just 2 +-- +-- >>> lowerBoundIn 2 5 xs 30 +-- Just 4 +-- +-- >>> lowerBoundIn 2 5 xs 40 +-- Just 4 +-- +-- >>> lowerBoundIn 2 5 xs 50 +-- Nothing +-- +-- @since 1.1.0.0 +{-# INLINE lowerBoundIn #-} +lowerBoundIn :: (VG.Vector v a, Ord a) => Int -> Int -> v a -> a -> Maybe Int +lowerBoundIn l_ r_ vec target + | ACIA.testInterval l r (VG.length vec) = bisectR l r $ \i -> VG.unsafeIndex vec i < target + | otherwise = Nothing + where + -- clamp + l = max 0 l_ + r = min (VG.length vec) r_ + +-- | \(O(\log n)\) Returns the index of the first element \(x\) in the vector such that +-- \(x \gt x_0\), or `Nothing` if no such element exists. +-- +-- @ +-- Y Y Y Y Y N N N N N Y: (<= x0) +-- --------- *---------> X N: (> x0) +-- R R: returning point +-- @ +-- +-- ==== __Example__ +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let xs = VU.fromList [10, 10, 20, 20, 40, 40] +-- >>> upperBound xs 10 +-- Just 2 +-- +-- >>> upperBound xs 20 +-- Just 4 +-- +-- >>> upperBound xs 30 +-- Just 4 +-- +-- >>> upperBound xs 40 +-- Nothing +-- +-- Out of range values: +-- +-- >>> upperBound xs 0 +-- Just 0 +-- +-- >>> upperBound xs 50 +-- Nothing +-- +-- @since 1.1.0.0 +{-# INLINE upperBound #-} +upperBound :: (HasCallStack, VG.Vector v a, Ord a) => v a -> a -> Maybe Int +upperBound vec = upperBoundIn 0 (VG.length vec) vec + +-- | \(O(\log n)\) Computes the `upperBound` for a slice of a vector within the interval \([l, r)\). +-- +-- - The user predicate evaluates indices in \([l, r)\). +-- - The interval \([l, r)\) is silently clamped to ensure it remains within the bounds \([0, n)\). +-- +-- ==== __Example__ +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let xs = VU.fromList [10, 10, 20, 20, 40, 40] +-- >>> -- *---*---* +-- >>> upperBoundIn 2 5 xs 0 +-- Just 2 +-- +-- >>> upperBoundIn 2 5 xs 10 +-- Just 2 +-- +-- >>> upperBoundIn 2 5 xs 20 +-- Just 4 +-- +-- >>> upperBoundIn 2 5 xs 30 +-- Just 4 +-- +-- >>> upperBoundIn 2 5 xs 40 +-- Nothing +-- +-- >>> upperBoundIn 2 5 xs 50 +-- Nothing +-- +-- @since 1.1.0.0 +{-# INLINE upperBoundIn #-} +upperBoundIn :: (VG.Vector v a, Ord a) => Int -> Int -> v a -> a -> Maybe Int +upperBoundIn l_ r_ vec target + | ACIA.testInterval l r (VG.length vec) = bisectR l r $ \i -> VG.unsafeIndex vec i <= target + | otherwise = Nothing + where + -- clamp + l = max 0 l_ + r = min (VG.length vec) r_ + +-- | \(O(\log n)\) Applies the bisection method on a half-open interval \([l, r)\) and returns the +-- left boundary point, or `Nothing` if no such point exists. +-- +-- @ +-- Y Y Y Y Y N N N N N Y: user predicate holds +-- --------* ----------> X N: user predicate does not hold +-- L L: the left boundary point returned +-- @ +-- +-- ==== __Example__ +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let xs = VU.fromList [10, 10, 20, 20, 30, 30] +-- >>> let n = VU.length xs +-- >>> bisectL 0 n ((<= 20) . (xs VU.!)) +-- Just 3 +-- +-- >>> bisectL 0 n ((<= 0) . (xs VU.!)) +-- Nothing +-- +-- >>> bisectL 0 n ((<= 100) . (xs VU.!)) +-- Just 5 +-- +-- >>> bisectL 0 3 ((<= 20) . (xs VU.!)) +-- Just 2 +-- +-- @since 1.1.0.0 +{-# INLINE bisectL #-} +bisectL :: (HasCallStack) => Int -> Int -> (Int -> Bool) -> Maybe Int +bisectL l r p = runIdentity $ bisectLM l r (pure . p) + +-- | \(O(\log n)\) Monadic variant of `bisectL`. +-- +-- @since 1.1.0.0 +{-# INLINE bisectLM #-} +bisectLM :: (HasCallStack, Monad m) => Int -> Int -> (Int -> m Bool) -> m (Maybe Int) +bisectLM l r p + | l >= r = pure Nothing + | otherwise = + bisectLImpl p l r <&> \case + i | i == (l - 1) -> Nothing + i -> Just i + +-- | \(O(\log n)\) Applies the bisection method on a half-open interval \([l, r)\) and returns the +-- right boundary point, or `Nothing` if no such point exists. +-- +-- +-- @ +-- Y Y Y Y Y N N N N N Y: user predicate holds +-- --------- *---------> X N: user predicate does not hold +-- R R: the right boundary point returned +-- @ +-- +-- ==== __Example__ +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let xs = VU.fromList [10, 10, 20, 20, 30, 30] +-- >>> let n = VU.length xs +-- >>> bisectR 0 n ((<= 20) . (xs VU.!)) +-- Just 4 +-- +-- >>> bisectR 0 n ((<= 0) . (xs VU.!)) +-- Just 0 +-- +-- >>> bisectR 0 n ((<= 100) . (xs VU.!)) +-- Nothing +-- +-- >>> bisectR 0 4 ((<= 20) . (xs VU.!)) +-- Nothing +-- +-- @since 1.1.0.0 +{-# INLINE bisectR #-} +bisectR :: (HasCallStack) => Int -> Int -> (Int -> Bool) -> Maybe Int +bisectR l r p = runIdentity $ bisectRM l r (pure . p) + +-- | \(O(\log n)\) Monadic variant of `bisectR`. +-- +-- @since 1.1.0.0 +{-# INLINE bisectRM #-} +bisectRM :: (HasCallStack, Monad m) => Int -> Int -> (Int -> m Bool) -> m (Maybe Int) +bisectRM l r p + | l >= r = pure Nothing + | otherwise = + bisectRImpl p l r <&> \case + i | i == r -> Nothing + i -> Just i diff --git a/src/AtCoder/Extra/Graph.hs b/src/AtCoder/Extra/Graph.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Graph.hs @@ -0,0 +1,147 @@ +-- | Re-export of the @Csr@ module and generic graph search functions. +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Graph + ( -- * Re-export of CSR + + -- | The `Csr.Csr` data type and all the functions such as `build` or `adj` are re-exported. + module Csr, + + -- * CSR helpers + swapDupe, + swapDupe', + scc, + + -- * Graph search + topSort, + ) +where + +import AtCoder.Extra.IntSet qualified as IS +import AtCoder.Internal.Buffer qualified as B +import AtCoder.Internal.Csr as Csr +import AtCoder.Internal.Scc qualified as ACISCC +import Control.Monad (when) +import Control.Monad.ST (runST) +import Data.Foldable (for_) +import Data.Vector qualified as V +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM + +-- | \(O(n)\) Converts non-directed edges into directional edges. This is a convenient function for +-- making an input to `build`. +-- +-- ==== __Example__ +-- `swapDupe` duplicates each edge reversing the direction: +-- +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> Gr.swapDupe $ VU.fromList [(0, 1, ()), (1, 2, ())] +-- [(0,1,()),(1,0,()),(1,2,()),(2,1,())] +-- +-- Create a non-directed graph: +-- +-- >>> let gr = Gr.build 3 . Gr.swapDupe $ VU.fromList [(0, 1, ()), (1, 2, ())] +-- >>> gr `Gr.adj` 0 +-- [1] +-- +-- >>> gr `Gr.adj` 1 +-- [0,2] +-- +-- >>> gr `Gr.adj` 2 +-- [1] +-- +-- @since 1.1.0.0 +{-# INLINE swapDupe #-} +swapDupe :: (VU.Unbox (Int, Int, w)) => VU.Vector (Int, Int, w) -> VU.Vector (Int, Int, w) +swapDupe = VU.concatMap (\(!u, !v, !w) -> VU.fromListN 2 [(u, v, w), (v, u, w)]) + +-- | \(O(n)\) Converts non-directed edges into directional edges. This is a convenient function for +-- making an input to `build'`. +-- +-- ==== __Example__ +-- `swapDupe'` duplicates each edge reversing the direction: +-- +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2)] +-- [(0,1),(1,0),(1,2),(2,1)] +-- +-- Create a non-directed graph: +-- +-- >>> let gr = Gr.build' 3 . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2)] +-- >>> gr `Gr.adj` 0 +-- [1] +-- +-- >>> gr `Gr.adj` 1 +-- [0,2] +-- +-- >>> gr `Gr.adj` 2 +-- [1] +-- +-- @since 1.1.0.0 +{-# INLINE swapDupe' #-} +swapDupe' :: (VU.Unbox (Int, Int)) => VU.Vector (Int, Int) -> VU.Vector (Int, Int) +swapDupe' = VU.concatMap (\(!u, !v) -> VU.fromListN 2 [(u, v), (v, u)]) + +-- | \(O(n + m)\) Returns the strongly connected components. +-- +-- ==== __Example__ +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> -- 0 == 1 -> 2 3 +-- >>> let gr = Gr.build' 4 $ VU.fromList [(0, 1), (1, 0), (1, 2)] +-- >>> Gr.scc gr +-- [[3],[0,1],[2]] +-- +-- @since 1.1.0.0 +{-# INLINE scc #-} +scc :: Csr w -> V.Vector (VU.Vector Int) +scc = ACISCC.sccCsr + +-- | \(O(n \log n + m)\) Returns the lexicographically smallest topological ordering of the given +-- graph. +-- +-- ==== Constraints +-- - The graph must be a DAG. +-- +-- ==== __Example__ +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let n = 5 +-- >>> let gr = Gr.build' n $ VU.fromList [(1, 2), (4, 0), (0, 3)] +-- >>> Gr.topSort n (gr `Gr.adj`) +-- [1,2,4,0,3] +-- +-- @since 1.1.0.0 +{-# INLINE topSort #-} +topSort :: Int -> (Int -> VU.Vector Int) -> VU.Vector Int +topSort n gr = runST $ do + inDeg <- VUM.replicate n (0 :: Int) + for_ [0 .. n - 1] $ \u -> do + VU.forM_ (gr u) $ \v -> do + VGM.modify inDeg (+ 1) v + + -- start from the vertices with zero in-degrees: + que <- IS.new n + inDeg' <- VU.unsafeFreeze inDeg + VU.iforM_ inDeg' $ \v d -> do + when (d == 0) $ do + IS.insert que v + + buf <- B.new n + let run = do + IS.deleteMin que >>= \case + Nothing -> pure () + Just u -> do + B.pushBack buf u + VU.forM_ (gr u) $ \v -> do + nv <- subtract 1 <$> VGM.read inDeg v + VGM.write inDeg v nv + when (nv == 0) $ do + IS.insert que v + run + + run + B.unsafeFreeze buf diff --git a/src/AtCoder/Extra/HashMap.hs b/src/AtCoder/Extra/HashMap.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/HashMap.hs @@ -0,0 +1,346 @@ +{-# LANGUAGE RecordWildCards #-} + +-- original implementaion: +-- + +-- | A dense, fast `Int` hash map with a fixed-sized `capacity` of \(n\). Most operations are +-- performed in \(O(1)\) time, but in average. +-- +-- NOTE: The entries (key - value pairs) cannot be invalidated due to the internal implementation +-- (called /open addressing/). +-- +-- ==== __Example__ +-- Create a `HashMap` with `capacity` \(10\): +-- +-- >>> import AtCoder.Extra.HashMap qualified as HM +-- >>> hm <- HM.new @_ @Int 10 +-- +-- `insert`, `lookup` and other functions are available in \(O(1)\) averaged time: +-- +-- >>> HM.insert hm 0 100 +-- >>> HM.insert hm 10 101 +-- >>> HM.size hm +-- 2 +-- +-- >>> HM.lookup hm 0 +-- Just 100 +-- +-- >>> HM.lookup hm 10 +-- Just 101 +-- +-- @since 1.1.0.0 +module AtCoder.Extra.HashMap + ( -- * HashMap + HashMap, + + -- * Constructors + new, + build, + + -- * Metadata + capacity, + size, + + -- * Lookups + lookup, + member, + notMember, + + -- * Modifications + + -- ** Insertions + insert, + insertWith, + exchange, + + -- ** Updates + modify, + modifyM, + + -- ** Reset + clear, + + -- * Conversions + + -- ** Safe conversions + keys, + elems, + assocs, + + -- ** Unsafe conversions + unsafeKeys, + unsafeElems, + unsafeAssocs, + ) +where + +import AtCoder.Internal.Assert qualified as ACIA +import Control.Monad (void, when) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Data.Bit (Bit (..)) +import Data.Bits (Bits (xor, (.&.)), (.>>.)) +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.Stack (HasCallStack) +import Prelude hiding (lookup) + +-- | A dense, fast `Int` hash map with a fixed-sized capacity of \(n\). +-- +-- @since 1.1.0.0 +data HashMap s a = HashMap + { -- | Maximum number of elements. + maxCapHM :: {-# UNPACK #-} !Int, + -- | The number of elements that can be added. + restCapHM :: !(VUM.MVector s Int), + -- | Bit mask for powerset iteration on indexing. + maskHM :: {-# UNPACK #-} !Int, + -- | Original key to the hash index. + keyHM :: !(VUM.MVector s Int), + -- | Values to the hash index. + valHM :: !(VUM.MVector s a), + -- | Whether the slot is used or not. + usedHM :: !(VUM.MVector s Bit) + } + +{-# INLINE decrementRestCapacity #-} +decrementRestCapacity :: (HasCallStack, PrimMonad m) => VUM.MVector (PrimState m) Int -> String -> m () +decrementRestCapacity restCap funcName = do + rest <- VGM.unsafeRead restCap 0 + let !_ = ACIA.runtimeAssert (rest > 0) $ "AtCoder.Extra.HashMap." ++ funcName ++ ": out of capacity" + VGM.unsafeWrite restCap 0 (rest - 1) + pure () + +-- | \(O(n)\) Creates a `HashMap` of capacity \(n\). +-- +-- @since 1.1.0.0 +{-# INLINE new #-} +new :: (PrimMonad m, VU.Unbox a) => Int -> m (HashMap (PrimState m) a) +new n = do + let !k0 = 1 + let !k = until (>= 2 * n) (* 2) k0 + let !maxCapHM = k `div` 2 + restCapHM <- VUM.replicate 1 maxCapHM + let !maskHM = k - 1 + keyHM <- VUM.unsafeNew k + valHM <- VUM.unsafeNew k + usedHM <- VUM.replicate k $ Bit False + pure HashMap {..} + +-- | \(O(n)\) Creates a `HashMap` of capacity \(n\) with initial entries. +-- +-- @since 1.1.0.0 +{-# INLINE build #-} +build :: (PrimMonad m, VU.Unbox a) => Int -> VU.Vector (Int, a) -> m (HashMap (PrimState m) a) +build n xs = do + hm <- new n + VU.forM_ xs $ \(!i, !x) -> do + insert hm i x + pure hm + +-- | \(O(1)\) Returns the maximum number of elements the hash map can store. +-- +-- @since 1.1.0.0 +{-# INLINE capacity #-} +capacity :: HashMap s a -> Int +capacity = maxCapHM + +-- | \(O(1)\) Returns the number of elements in the hash map. +-- +-- @since 1.1.0.0 +{-# INLINE size #-} +size :: (PrimMonad m) => HashMap (PrimState m) a -> m Int +size HashMap {..} = do + !rest <- VUM.unsafeRead restCapHM 0 + pure $ maxCapHM - rest + +-- | \(O(1)\) (Internal) Hash value calculation. +-- +-- @since 1.1.0.0 +{-# INLINE hash #-} +hash :: HashMap a s -> Int -> Int +hash hm x = fromIntegral $ (x3 `xor` (x3 .>>. 31)) .&. fromIntegral (maskHM hm) + where + fixedRandom, x1, x2, x3 :: Word64 + fixedRandom = 321896566547 + x1 = fromIntegral x + fixedRandom + x2 = (x1 `xor` (x1 .>>. 30)) * 0xbf58476d1ce4e5b9 + x3 = (x2 `xor` (x2 .>>. 27)) * 0x94d049bb133111eb + +-- | \(O(1)\) (Internal) Hashed slot search. +-- +-- @since 1.1.0.0 +{-# INLINE index #-} +index :: (PrimMonad m) => HashMap (PrimState m) a -> Int -> m Int +index hm@HashMap {..} k = inner (hash hm k) + where + inner !h = do + Bit b <- VGM.read usedHM h + -- already there? + k' <- VGM.read keyHM h + if b && k' /= k + then inner $ (h + 1) .&. maskHM + else pure h + +-- | \(O(1)\) Return the value to which the specified key is mapped, or `Nothing` if this map +-- contains no mapping for the key. +-- +-- @since 1.1.0.0 +{-# INLINE lookup #-} +lookup :: (HasCallStack, VU.Unbox a, PrimMonad m) => HashMap (PrimState m) a -> Int -> m (Maybe a) +lookup hm@HashMap {..} k = do + i <- index hm k + Bit b <- VGM.read usedHM i + if b + then Just <$> VGM.read valHM i + else pure Nothing + +-- | \(O(1)\) Checks whether the hash map contains the element. +-- +-- @since 1.1.0.0 +{-# INLINE member #-} +member :: (HasCallStack, PrimMonad m) => HashMap (PrimState m) a -> Int -> m Bool +member hm@HashMap {..} k = do + i <- index hm k + Bit b <- VGM.read usedHM i + -- TODO: is this key check necessary + k' <- VGM.read keyHM i + pure $ b && k' == k + +-- | \(O(1)\) Checks whether the hash map does not contain the element. +-- +-- @since 1.1.0.0 +{-# INLINE notMember #-} +notMember :: (HasCallStack, PrimMonad m) => HashMap (PrimState m) a -> Int -> m Bool +notMember hm k = not <$> member hm k + +-- | \(O(1)\) Inserts a \((k, v)\) pair. +-- +-- @since 1.1.0.0 +{-# INLINE insert #-} +insert :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> Int -> a -> m () +insert hm k v = void $ exchange hm k v + +-- | \(O(1)\) Inserts a \((k, v)\) pair. If the key exists, the function will insert the pair +-- \((k, f(v_{\mathrm{new}}, v_{\mathrm{old}}))\). +-- +-- @since 1.1.0.0 +{-# INLINE insertWith #-} +insertWith :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m () +insertWith hm@HashMap {..} f k v = do + i <- index hm k + Bit b <- VGM.exchange usedHM i $ Bit True + if b + then do + -- modify the existing entry + VGM.modify valHM (f v) i + else do + -- insert the new \((k, v)\) pair + decrementRestCapacity restCapHM "insertWith" + VGM.write keyHM i k + VGM.write valHM i v + +-- | \(O(1)\) Inserts a \((k, v)\) pair and returns the old value, or `Nothing` if no such entry +-- exists. +-- +-- @since 1.1.0.0 +{-# INLINE exchange #-} +exchange :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> Int -> a -> m (Maybe a) +exchange hm@HashMap {..} k v = do + i <- index hm k + Bit b <- VGM.exchange usedHM i $ Bit True + if b + then do + -- overwrite the existing entry + Just <$> VGM.exchange valHM i v + else do + -- insert the new (key, value) pair + decrementRestCapacity restCapHM "exchange" + VGM.write keyHM i k + VGM.write valHM i v + pure Nothing + +-- | \(O(1)\) Modifies the element at the given key. Does nothing if no such entry exists. +-- +-- @since 1.1.0.0 +{-# INLINE modify #-} +modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> a) -> Int -> m () +modify hm@HashMap {..} f k = do + i <- index hm k + Bit b <- VGM.read usedHM i + when b $ do + VGM.modify valHM f i + +-- | \(O(1)\) Modifies the element at the given key. Does nothing if no such entry exists. +-- +-- @since 1.1.0.0 +{-# INLINE modifyM #-} +modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> m a) -> Int -> m () +modifyM hm@HashMap {..} f k = do + i <- index hm k + Bit b <- VGM.read usedHM i + when b $ do + VGM.modifyM valHM f i + +-- | \(O(n)\) Clears all the elements. +-- +-- @since 1.1.0.0 +{-# INLINE clear #-} +clear :: (PrimMonad m) => HashMap (PrimState m) a -> m () +clear HashMap {..} = do + VGM.set usedHM $ Bit False + VUM.unsafeWrite restCapHM 0 maxCapHM + +-- | \(O(n)\) Enumerates the keys in the hash map. +-- +-- @since 1.1.0.0 +{-# INLINE keys #-} +keys :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector Int) +keys hm = VU.force <$> unsafeKeys hm + +-- | \(O(n)\) Enumerates the elements (values) in the hash map. +-- +-- @since 1.1.0.0 +{-# INLINE elems #-} +elems :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector a) +elems hm = VU.force <$> unsafeElems hm + +-- | \(O(n)\) Enumerates the key-value pairs in the hash map. +-- +-- @since 1.1.0.0 +{-# INLINE assocs #-} +assocs :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector (Int, a)) +assocs hm = VU.force <$> unsafeAssocs hm + +-- | \(O(n)\) Enumerates the keys in the hash map. +-- +-- @since 1.1.0.0 +{-# INLINE unsafeKeys #-} +unsafeKeys :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector Int) +unsafeKeys HashMap {..} = do + used <- VU.unsafeFreeze usedHM + keys_ <- VU.unsafeFreeze keyHM + pure $ VU.ifilter (const . unBit . (used VG.!)) keys_ + +-- | \(O(n)\) Enumerates the elements (values) in the hash map. +-- +-- @since 1.1.0.0 +{-# INLINE unsafeElems #-} +unsafeElems :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector a) +unsafeElems HashMap {..} = do + used <- VU.unsafeFreeze usedHM + vals <- VU.unsafeFreeze valHM + pure $ VU.ifilter (const . unBit . (used VG.!)) vals + +-- | \(O(n)\) Enumerates the key-value pairs in the hash map. +-- +-- @since 1.1.0.0 +{-# INLINE unsafeAssocs #-} +unsafeAssocs :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector (Int, a)) +unsafeAssocs HashMap {..} = do + used <- VU.unsafeFreeze usedHM + keys_ <- VU.unsafeFreeze keyHM + vals <- VU.unsafeFreeze valHM + pure $ VU.ifilter (const . unBit . (used VG.!)) $ VU.zip keys_ vals diff --git a/src/AtCoder/Extra/IntMap.hs b/src/AtCoder/Extra/IntMap.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/IntMap.hs @@ -0,0 +1,338 @@ +{-# LANGUAGE RecordWildCards #-} + +-- | A dense, fast `Int` map implemented as a 64-ary tree that covers an interval \([0, n)\). +-- +-- ==== __Example__ +-- Create an `IntMap` with capacity \(10\): +-- +-- >>> import AtCoder.Extra.IntMap qualified as IM +-- >>> im <- IM.new @_ @Int 10 +-- +-- `insert`, `delete`, `lookup` and other functions are available: +-- +-- >>> IM.insert im 0 100 +-- >>> IM.insert im 9 101 +-- >>> IM.delete im 0 +-- True +-- +-- >>> IM.size im +-- 1 +-- +-- >>> IM.lookup im 9 +-- Just 101 +-- +-- >>> IM.lookup im 1 +-- Nothing +-- +-- >>> IM.lookupGT im 5 +-- Just (9,101) +-- +-- @since 1.1.0.0 +module AtCoder.Extra.IntMap + ( -- * IntMap + IntMap, + + -- * Constructors + new, + build, + + -- * Metadata + capacity, + size, + null, + + -- * Lookups + lookup, + member, + notMember, + + -- ** Compartive lookups + lookupGE, + lookupGT, + lookupLE, + lookupLT, + + -- ** Max/Min lookups + lookupMin, + lookupMax, + + -- * Modifications + + -- ** Insertions + insert, + insertWith, + + -- ** Updates + modify, + modifyM, + + -- ** Deletions + delete, + delete_, + deleteMin, + deleteMax, + + -- * Conversions + keys, + elems, + assocs, + ) +where + +import AtCoder.Extra.IntSet qualified as IS +import Control.Monad (when) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Data.Maybe (fromJust) +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 Prelude hiding (lookup, null) + +-- | A dense, fast `Int` map implemented as a 64-ary tree that covers an interval \([0, n)\). +-- +-- @since 1.1.0.0 +data IntMap s a = IntMap + { setIM :: !(IS.IntSet s), + valIM :: !(VUM.MVector s a) + } + +-- | \(O(n)\) Creates an `IntMap` for an interval \([0, n)\). +-- +-- @since 1.1.0.0 +{-# INLINE new #-} +new :: (PrimMonad m, VU.Unbox a) => Int -> m (IntMap (PrimState m) a) +new cap = do + setIM <- IS.new cap + valIM <- VUM.unsafeNew cap + pure IntMap {..} + +-- | \(O(n + m \log n)\) Creates an `IntMap` for an interval \([0, n)\) with initial values. +-- +-- @since 1.1.0.0 +{-# INLINE build #-} +build :: (PrimMonad m, VU.Unbox a) => Int -> VU.Vector (Int, a) -> m (IntMap (PrimState m) a) +build cap xs = do + im <- new cap + VU.forM_ xs $ \(!k, !v) -> do + insert im k v + pure im + +-- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is covered by the map. +-- +-- @since 1.1.0.0 +{-# INLINE capacity #-} +capacity :: IntMap s a -> Int +capacity = IS.capacity . setIM + +-- | \(O(1)\) Returns the number of elements in the map. +-- +-- @since 1.1.0.0 +{-# INLINE size #-} +size :: (PrimMonad m) => IntMap (PrimState m) a -> m Int +size = IS.size . setIM + +-- | \(O(1)\) Returns whether the map is empty. +-- +-- @since 1.1.0.0 +{-# INLINE null #-} +null :: (PrimMonad m) => IntMap (PrimState m) a -> m Bool +null = IS.null . setIM + +-- | \(O(\log n)\) Looks up the value for a key. +-- +-- @since 1.1.0.0 +{-# INLINE lookup #-} +lookup :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe a) +lookup im@IntMap {..} k = do + member im k >>= \case + True -> Just <$> VGM.read valIM k + False -> pure Nothing + +-- | \(O(\log n)\) Tests whether a key \(k\) is in the map. +-- +-- @since 1.1.0.0 +{-# INLINE member #-} +member :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool +member = IS.member . setIM + +-- | \(O(\log n)\) Tests whether a key \(k\) is not in the map. +-- +-- @since 1.1.0.0 +{-# INLINE notMember #-} +notMember :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool +notMember = IS.notMember . setIM + +-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the smallest key \(k\) such that \(k \ge k_0\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupGE #-} +lookupGE :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a)) +lookupGE IntMap {..} k = do + IS.lookupGE setIM k >>= \case + Just i -> Just . (i,) <$> VGM.read valIM i + Nothing -> pure Nothing + +-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the smallest \(k\) such that \(k \gt k_0\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupGT #-} +lookupGT :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a)) +lookupGT is k = lookupGE is (k + 1) + +-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the largest key \(k\) such that \(k \le k_0\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupLE #-} +lookupLE :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a)) +lookupLE IntMap {..} k = do + IS.lookupLE setIM k >>= \case + Just i -> Just . (i,) <$> VGM.read valIM i + Nothing -> pure Nothing + +-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the largest key \(k\) such that \(k \lt k_0\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupLT #-} +lookupLT :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a)) +lookupLT is k = lookupLE is (k - 1) + +-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the minimum key \(k\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupMin #-} +lookupMin :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a)) +lookupMin is = lookupGE is 0 + +-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the maximum key \(k\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupMax #-} +lookupMax :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a)) +lookupMax im = lookupLE im (IS.capacity (setIM im) - 1) + +-- | \(O(\log n)\) Inserts a \((k, v)\) pair into the map. If an entry with the same key already +-- exists, it is overwritten. +-- +-- @since 1.1.0.0 +{-# INLINE insert #-} +insert :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> a -> m () +insert IntMap {..} k v = do + IS.insert setIM k + VGM.write valIM k v + +-- | \(O(\log n)\) Inserts a \((k, v)\) pair into the map. If an entry with the same key already +-- exists, it overwritten with \(f(v_{\mathrm{new}}, v_{\mathrm{old}})\). +-- +-- @since 1.1.0.0 +{-# INLINE insertWith #-} +insertWith :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m () +insertWith IntMap {..} f k v = do + b <- IS.member setIM k + if b + then do + VGM.modify valIM (f v) k + else do + IS.insert setIM k + VGM.write valIM k v + +-- | \(O(\log n)\) Modifies the value associated with a key. If an entry with the same key already +-- does not exist, nothing is performed. +-- +-- @since 1.1.0.0 +{-# INLINE modify #-} +modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> a) -> Int -> m () +modify IntMap {..} f k = do + b <- IS.member setIM k + when b $ do + VGM.modify valIM f k + +-- | \(O(\log n)\) Modifies the value associated with a key. If an entry with the same key already +-- does not exist, nothing is performed. +-- +-- @since 1.1.0.0 +{-# INLINE modifyM #-} +modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> m a) -> Int -> m () +modifyM IntMap {..} f k = do + b <- IS.member setIM k + when b $ do + VGM.modifyM valIM f k + +-- | \(O(\log n)\) Deletes the \((k, v)\) pair with the key \(k\) from the map. Does nothing if no +-- such key exists. Returns whether the key existed. +-- +-- @since 1.1.0.0 +{-# INLINE delete #-} +delete :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool +delete im = IS.delete (setIM im) + +-- | \(O(\log n)\) Deletes the \((k, v)\) pair with the key \(k\) from the map. Does nothing if no +-- such key exists. +-- +-- @since 1.1.0.0 +{-# INLINE delete_ #-} +delete_ :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m () +delete_ im = IS.delete_ (setIM im) + +-- | \(O(\log n)\) Deletes the \((k, v)\) pair with the minimum key \(k\) in the map. +-- +-- @since 1.1.0.0 +{-# INLINE deleteMin #-} +deleteMin :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a)) +deleteMin is = do + lookupMin is + >>= mapM + ( \(!key, !val) -> do + delete_ is key + pure (key, val) + ) + +-- | \(O(\log n)\) Deletes the \((k, v)\) pair with maximum key \(k\) in the map. +-- +-- @since 1.1.0.0 +{-# INLINE deleteMax #-} +deleteMax :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a)) +deleteMax is = do + lookupMax is + >>= mapM + ( \(!k, !v) -> do + delete_ is k + pure (k, v) + ) + +-- | \(O(n \log n)\) Enumerates the keys in the map. +-- +-- @since 1.1.0.0 +{-# INLINE keys #-} +keys :: (PrimMonad m) => IntMap (PrimState m) a -> m (VU.Vector Int) +keys = IS.keys . setIM + +-- | \(O(n \log n)\) Enumerates the elements (values) in the map. +-- +-- @since 1.1.0.0 +{-# INLINE elems #-} +elems :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (VU.Vector a) +elems im@IntMap {..} = do + n <- IS.size setIM + VU.unfoldrExactNM + n + ( \i -> do + (!i', !x') <- fromJust <$> lookupGT im i + pure (x', i') + ) + (-1) + +-- | \(O(n \log n)\) Enumerates the key-value pairs in the map. +-- +-- @since 1.1.0.0 +{-# INLINE assocs #-} +assocs :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (VU.Vector (Int, a)) +assocs im@IntMap {..} = do + n <- IS.size setIM + VU.unfoldrExactNM + n + ( \i -> do + (!i', !x') <- fromJust <$> lookupGT im i + pure ((i', x'), i') + ) + (-1) diff --git a/src/AtCoder/Extra/IntSet.hs b/src/AtCoder/Extra/IntSet.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/IntSet.hs @@ -0,0 +1,391 @@ +{-# LANGUAGE RecordWildCards #-} + +-- original implementation: +-- + +-- | A dense, fast `Int` set implemented as a 64-ary tree that covers an interval \([0, n)\). +-- +-- ==== __Example__ +-- Create an `IntSet` with capacity \(10\): +-- +-- >>> import AtCoder.Extra.IntSet qualified as IS +-- >>> is <- IS.new @_ 10 +-- +-- `insert`, `delete` and other functions are available: +-- +-- >>> IS.insert is 0 +-- >>> IS.insert is 9 +-- >>> IS.member is 9 +-- True +-- +-- >>> IS.delete is 0 +-- True +-- +-- >>> IS.size is +-- 1 +-- +-- >>> IS.member is 1 +-- False +-- +-- >>> IS.lookupGT is 5 +-- Just 9 +-- +-- @since 1.1.0.0 +module AtCoder.Extra.IntSet + ( -- * IntSet + IntSet, + + -- * Constructors + new, + build, + + -- * Metadata + capacity, + size, + null, + + -- * Lookups + member, + notMember, + + -- ** Compartive lookups + lookupGE, + lookupGT, + lookupLE, + lookupLT, + + -- ** Max/Min lookups + lookupMin, + lookupMax, + + -- * Modifications + + -- ** Insertions + insert, + + -- ** Deletions + delete, + delete_, + deleteMin, + deleteMax, + + -- * Conversions + keys, + ) +where + +import AtCoder.Internal.Assert qualified as ACIA +import Control.Monad (unless, void) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Data.Bifunctor (bimap) +import Data.Bits + ( Bits (clearBit, setBit, testBit), + FiniteBits (countLeadingZeros, countTrailingZeros), + (.<<.), + (.>>.), + ) +import Data.Maybe (fromJust) +import Data.Vector qualified as V +import Data.Vector.Generic qualified as VG +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import GHC.Stack (HasCallStack) +import Prelude hiding (null) + +-- | \(O(1)\) Retrieves the most significant bit. +-- +-- >>> msbOf 0 +-- -1 +-- +-- >>> msbOf maxBound +-- 62 +-- +-- >>> msbOf $ 4 + 2 + 1 +-- 2 +{-# INLINE msbOf #-} +msbOf :: Int -> Int +msbOf !x = 63 - countLeadingZeros x + +-- | \(O(1)\) Retrieves the least significant bit. +-- +-- >>> lsbOf 0 +-- -1 +-- +-- >>> lsbOf maxBound +-- 0 +-- +-- >>> lsbOf $ 4 + 2 + 1 +-- 0 +{-# INLINE lsbOf #-} +lsbOf :: Int -> Int +lsbOf 0 = -1 +lsbOf x = countTrailingZeros x + +{-# INLINE wordSize #-} +wordSize :: Int +wordSize = 64 + +-- | A dense, fast `Int` set implemented as a 64-ary tree that covers an interval \([0, n)\). +-- +-- @since 1.1.0.0 +data IntSet s = IntSet + { -- | Maximum number of elements. + capacityIS :: {-# UNPACK #-} !Int, + -- | The number of elements. + sizeIS :: !(VUM.MVector s Int), + -- | Segments. + vecIS :: !(V.Vector (VUM.MVector s Int)) + } + +-- | \(O(n)\) Creates an `IntSet` for the interval \([0, n)\). +-- +-- @since 1.1.0.0 +{-# INLINE new #-} +new :: (PrimMonad m) => Int -> m (IntSet (PrimState m)) +new capacityIS = do + vecIS <- + V.unfoldrExactNM + (max 1 logSize) + ( \len -> do + let !len' = (len + wordSize - 1) `div` wordSize + (,len') <$> VUM.replicate len' 0 + ) + capacityIS + sizeIS <- VUM.replicate 1 (0 :: Int) + pure IntSet {..} + where + (!_, !logSize) = + until + ((<= 1) . fst) + (bimap ((`div` wordSize) . (+ (wordSize - 1))) (+ 1)) + (capacityIS, 0) + +-- | \(O(n + m \log n)\) Creates an `IntSet` for the interval \([0, n)\) with initial values. +-- +-- @since 1.1.0.0 +{-# INLINE build #-} +build :: (PrimMonad m) => Int -> VU.Vector Int -> m (IntSet (PrimState m)) +build n vs = do + set <- new n + VU.forM_ vs (insert set) + pure set + +-- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is covered by the set. +-- +-- @since 1.1.0.0 +{-# INLINE capacity #-} +capacity :: IntSet s -> Int +capacity = capacityIS + +-- | \(O(1)\) Returns the number of elements in the set. +-- +-- @since 1.1.0.0 +{-# INLINE size #-} +size :: (PrimMonad m) => IntSet (PrimState m) -> m Int +size = (`VUM.unsafeRead` 0) . sizeIS + +-- | \(O(1)\) Returns whether the set is empty. +-- +-- @since 1.1.0.0 +{-# INLINE null #-} +null :: (PrimMonad m) => IntSet (PrimState m) -> m Bool +null = ((== 0) <$>) . size + +-- | \(O(\log n)\) Tests whether \(k\) is in the set. +-- +-- @since 1.1.0.0 +{-# INLINE member #-} +member :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool +member IntSet {..} k + | ACIA.testIndex k capacityIS = do + let (!q, !r) = k `divMod` wordSize + (`testBit` r) <$> VGM.unsafeRead (VG.unsafeHead vecIS) q + | otherwise = pure False + +-- | \(O(\log n)\) Tests whether \(k\) is not in the set. +-- +-- @since 1.1.0.0 +{-# INLINE notMember #-} +notMember :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool +notMember dis k = not <$> member dis k + +-- | \(O(\log n)\) Looks up the smallest key \(k\) such that \(k \ge k_0\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupGE #-} +lookupGE :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int) +lookupGE IntSet {..} i0 + | i0 >= capacityIS = pure Nothing + | otherwise = inner 0 $ max 0 i0 -- REMARK: it's very important to keep @i@ non-negative. + where + inner h i + | h >= V.length vecIS = pure Nothing + -- ? + | q == VUM.length (VG.unsafeIndex vecIS h) = pure Nothing + | otherwise = do + d <- (.>>. r) <$> VGM.unsafeRead (VG.unsafeIndex vecIS h) q + if d == 0 + then inner (h + 1) (q + 1) + else + Just + <$> V.foldM' + ( \ !acc vec -> do + !dx <- lsbOf <$> VGM.unsafeRead vec acc + pure $ acc * wordSize + dx + ) + (i + lsbOf d) + (V.unsafeBackpermute vecIS (V.enumFromStepN (h - 1) (-1) h)) + where + (!q, !r) = i `divMod` wordSize + +-- | \(O(\log n)\) Looks up the smallest key \(k\) such that \(k \gt k_0\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupGT #-} +lookupGT :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int) +lookupGT is k = lookupGE is (k + 1) + +-- | \(O(\log n)\) Looks up the largest key \(k\) such that \(k \le k_0\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupLE #-} +lookupLE :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int) +lookupLE IntSet {..} i0 + | i0 <= -1 = pure Nothing + | otherwise = inner 0 $ min (capacityIS - 1) i0 + where + inner h i + | h >= V.length vecIS = pure Nothing + | i == -1 = pure Nothing + | otherwise = do + d <- (.<<. (63 - r)) <$> VGM.unsafeRead (VG.unsafeIndex vecIS h) q + if d == 0 + then inner (h + 1) (q - 1) + else do + Just + <$> V.foldM' + ( \ !acc vec -> do + !dx <- msbOf <$> VGM.unsafeRead vec acc + pure $ acc * wordSize + dx + ) + (i - countLeadingZeros d) + (V.unsafeBackpermute vecIS (V.enumFromStepN (h - 1) (-1) h)) + where + (!q, !r) = i `divMod` wordSize + +-- | \(O(\log n)\) Looks up the largest key \(k\) such that \(k \lt k_0\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupLT #-} +lookupLT :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int) +lookupLT is k = lookupLE is (k - 1) + +-- | \(O(\log n)\) Looks up the minimum key. +-- +-- @since 1.1.0.0 +{-# INLINE lookupMin #-} +lookupMin :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int) +lookupMin is = lookupGE is 0 + +-- | \(O(\log n)\) Looks up the maximum key. +-- +-- @since 1.1.0.0 +{-# INLINE lookupMax #-} +lookupMax :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int) +lookupMax is = lookupLE is (capacityIS is - 1) + +-- | \(O(\log n)\) Inserts a key \(k\) into the set. If an entry with the same key already exists, +-- it is overwritten. +-- +-- @since 1.1.0.0 +{-# INLINE insert #-} +insert :: (HasCallStack, PrimMonad m) => IntSet (PrimState m) -> Int -> m () +insert is@IntSet {..} k = do + b <- member is k + unless b $ do + VUM.unsafeModify sizeIS (+ 1) 0 + V.foldM'_ + ( \i vec -> do + let (!q, !r) = i `divMod` wordSize + VGM.unsafeModify vec (`setBit` r) q + pure q + ) + k + vecIS + where + !_ = ACIA.checkIndex "AtCoder.Extra.IntSet.insert" k capacityIS + +-- | \(O(\log n)\) Deletes a key \(k\) from the set. Does nothing if no such key exists. Returns +-- whether the key existed. +-- +-- @since 1.1.0.0 +{-# INLINE delete #-} +delete :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool +delete is@IntSet {..} k = do + b_ <- member is k + if b_ + then do + VUM.unsafeModify sizeIS (subtract 1) 0 + V.foldM'_ + ( \(!b, !i) vec -> do + let (!q, !r) = i `divMod` wordSize + -- TODO: early return is possible + unless b $ do + VGM.unsafeModify vec (`clearBit` r) q + -- `b` remembers if any other bit was on + b' <- (/= 0) <$> VGM.unsafeRead vec q + pure (b', q) + ) + (False, k) + vecIS + pure True + else pure False + +-- | \(O(\log n)\) Deletes a key \(k\) from the set. Does nothing if no such key exists. +-- +-- @since 1.1.0.0 +{-# INLINE delete_ #-} +delete_ :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m () +delete_ is k = void $ delete is k + +-- | \(O(\log n)\) Deletes the minimum key from the set. Returns `Nothing` if the set is empty. +-- +-- @since 1.1.0.0 +{-# INLINE deleteMin #-} +deleteMin :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int) +deleteMin is = do + lookupMin is + >>= mapM + ( \key -> do + delete_ is key + pure key + ) + +-- | \(O(\log n)\) Deletes the maximum key from the set. Returns `Nothing` if the set is empty. +-- +-- @since 1.1.0.0 +{-# INLINE deleteMax #-} +deleteMax :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int) +deleteMax is = do + lookupMax is + >>= mapM + ( \key -> do + delete_ is key + pure key + ) + +-- | \(O(n \log n)\) Enumerates the keys in the map. +-- +-- @since 1.1.0.0 +{-# INLINE keys #-} +keys :: (PrimMonad m) => IntSet (PrimState m) -> m (VU.Vector Int) +keys is@IntSet {sizeIS} = do + n <- VGM.unsafeRead sizeIS 0 + VU.unfoldrExactNM + n + ( \i -> do + i' <- fromJust <$> lookupGT is i + pure (i', i') + ) + (-1) diff --git a/src/AtCoder/Extra/IntervalMap.hs b/src/AtCoder/Extra/IntervalMap.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/IntervalMap.hs @@ -0,0 +1,462 @@ +{-# LANGUAGE DerivingStrategies #-} + +-- original implementation: +-- + +-- | Dense map covering \([0, n)\) that manages non-overlapping intervals \([l, r)\) within it. Each +-- interval has an associated value \(v\). Use @onAdd@ and @onDel@ hooks to track interval state +-- changes during `buildM`, `insertM` and `deleteM` operations. +-- +-- ==== Invariant +-- Each interval is operated as a whole, similar to a persistant data structure. When part of an +-- inerval is modified, the whole interval is deleted first, and the subintervals are re-inserted. +-- It's important for tracking non-linear interval information with the @onAdd@ and @onDel@ hooks +-- (callbacks). +-- +-- ==== __Example__ +-- Create an `IntervalMap` that covers a half-open interval \([0, n)\): +-- +-- >>> import AtCoder.Extra.IntervalMap qualified as ITM +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> import Data.Vector.Unboxed.Mutable qualified as VUM +-- >>> itm <- ITM.new @_ @Int 4 +-- +-- It handles range set queries in amortized \(O(\log n)\) time: +-- +-- >>> ITM.insert itm 0 4 0 -- 0 0 0 0 +-- >>> ITM.insert itm 1 3 1 -- 0 1 1 0 +-- >>> ITM.freeze itm +-- [(0,(1,0)),(1,(3,1)),(3,(4,0))] +-- +-- Track interval informations with the @onAdd@ and @onDel@ hooks: +-- +-- >>> import Debug.Trace (traceShow) +-- >>> itm <- ITM.new @_ @Int 4 +-- >>> let onAdd l r x = print ("onAdd", l, r, x) +-- >>> let onDel l r x = print ("onDel", l, r, x) +-- +-- >>> ITM.insertM itm 0 4 0 onAdd onDel -- 0 0 0 0 +-- ("onAdd",0,4,0) +-- +-- >>> ITM.insertM itm 1 3 1 onAdd onDel -- 0 1 1 0 +-- ("onDel",0,4,0) +-- ("onAdd",0,1,0) +-- ("onAdd",3,4,0) +-- ("onAdd",1,3,1) +-- +-- >>> ITM.deleteM itm 0 4 onAdd onDel +-- ("onDel",0,1,0) +-- ("onDel",1,3,1) +-- ("onDel",3,4,0) +-- +-- @since 1.1.0.0 +module AtCoder.Extra.IntervalMap + ( -- * IntervalMap + IntervalMap, + + -- * Constructors + new, + build, + buildM, + + -- * Metadata + capacity, + + -- * Lookups + contains, + intersects, + lookup, + read, + readMaybe, + + -- * Modifications + + -- ** Insertions + insert, + insertM, + + -- ** Deletions + delete, + deleteM, + + -- ** Overwrites + overwrite, + overwriteM, + + -- * Conversions + freeze, + ) +where + +import AtCoder.Extra.IntMap qualified as IM +import Control.Monad (foldM_) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Data.Vector.Generic qualified as G +import Data.Vector.Unboxed qualified as VU +import GHC.Stack (HasCallStack) +import Prelude hiding (lookup, read) + +-- | Dense map covering \([0, n)\) that manages non-overlapping intervals \((l, r)\) within it. Each +-- interval has an associated value \(x\). Use @onAdd@ and @onDel@ hooks to track interval state +-- changes during `buildM`, `insertM` and `deleteM` operations. +-- +-- @since 1.1.0.0 +newtype IntervalMap s a = IntervalMap + { -- | Maps \(l\) to \((r, a)\). + unITM :: IM.IntMap s (Int, a) + } + +-- | \(O(n)\) Creates an empty `IntervalMap`. +-- +-- @since 1.1.0.0 +new :: (PrimMonad m, VU.Unbox a) => Int -> m (IntervalMap (PrimState m) a) +new = fmap IntervalMap . IM.new + +-- | \(O(n + m \log n)\) Creates an `IntervalMap` by combining consecutive equal values into one +-- interval. +-- +-- ==== __Example__ +-- >>> itm <- build @_ @Int (VU.fromList [10,10,11,11,12,12]) +-- >>> freeze itm +-- [(0,(2,10)),(2,(4,11)),(4,(6,12))] +-- +-- @since 1.1.0.0 +build :: (PrimMonad m, Eq a, VU.Unbox a) => VU.Vector a -> m (IntervalMap (PrimState m) a) +build xs = buildM xs onAdd + where + onAdd _ _ _ = pure () + +-- | \(O(n + m \log n)\) Creates an `IntervalMap` by combining consecutive equal values into one +-- interval, while performing @onAdd@ hook for each interval. +-- +-- @since 1.1.0.0 +buildM :: + (PrimMonad m, Eq a, VU.Unbox a) => + -- | Input values + VU.Vector a -> + -- | @onAdd@ hook that take an interval \([l, r)\) with associated value \(v\) + (Int -> Int -> a -> m ()) -> + -- | The map + m (IntervalMap (PrimState m) a) +buildM xs onAdd = do + dim <- IM.new (G.length xs) + foldM_ (step dim) (0 :: Int) $ G.group xs + pure $ IntervalMap dim + where + step dim !l !xs' = do + let !l' = l + G.length xs' + IM.insert dim l (l', G.head xs') + onAdd l l' (G.head xs') + pure l' + +-- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is managed by the map. +-- +-- @since 1.1.0.0 +{-# INLINE capacity #-} +capacity :: IntervalMap s a -> Int +capacity = IM.capacity . unITM + +-- | \(O(\log n)\) Returns whether a point \(x\) is contained within any of the intervals. +-- +-- @since 1.1.0.0 +contains :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> m Bool +contains itm i = intersects itm i (i + 1) + +-- | \(O(\log n)\) Returns whether an interval \([l, r)\) is fully contained within any of the +-- intervals. +-- +-- @since 1.1.0.0 +intersects :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m Bool +intersects (IntervalMap dim) l r + | l >= r = pure False + | otherwise = do + res <- IM.lookupLE dim l + pure $ case res of + Just (!_, (!r', !_)) -> r <= r' + _ -> False + +-- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\). +-- +-- @since 1.1.0.0 +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 + | otherwise = do + res <- IM.lookupLE im l + pure $ case res of + Just (!l', (!r', !a)) + | r <= r' -> Just (l', r', a) + _ -> Nothing + +-- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\) and reads out the value. +-- Throws an error if no such interval exists. +-- +-- @since 1.1.0.0 +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 + pure $ case res of + Just !a -> a + Nothing -> error $ "[read] not a member: " ++ show (l, r) + +-- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\) and reads out the value. +-- Returns `Nothing` if no such interval exists. +-- +-- @since 1.1.0.0 +readMaybe :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m (Maybe a) +readMaybe (IntervalMap dim) l r + | l >= r = pure Nothing + | otherwise = do + res <- IM.lookupLE dim l + pure $ case res of + Just (!_, (!r', !a)) + | r <= r' -> Just a + _ -> Nothing + +-- | Amortized \(O(\log n)\) Inserts an interval \([l, r)\) with associated value \(v\) into the +-- map. Overwrites any overlapping intervals. +-- +-- @since 1.1.0.0 +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 + onAdd _ _ _ = pure () + onDel _ _ _ = pure () + +-- | Amortized \(O(\log n)\) Inserts an interval \([l, r)\) with associated value \(v\) into the +-- map. Overwrites any overlapping intervals. Tracks interval state changes via @onAdd@ and @onDel@ +-- hooks. +-- +-- @since 1.1.0.0 +insertM :: + (PrimMonad m, Eq a, VU.Unbox a) => + -- | The map + IntervalMap (PrimState m) a -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(v\) + a -> + -- | @onAdd@ hook that take an interval \([l, r)\) with associated value \(v\) + (Int -> Int -> a -> m ()) -> + -- | @onDel@ hook that take an interval \([l, r)\) with associated value \(v\) + (Int -> Int -> a -> m ()) -> + m () +insertM (IntervalMap dim) l0 r0 x onAdd onDel + | l0 >= r0 = pure () + | otherwise = do + !r <- handleRight l0 r0 + (!l', !r') <- handleLeft l0 r + onAdd l' r' x + IM.insert dim l' (r', x) + where + handleRight l r = do + res <- IM.lookupGE dim l + case res of + Just interval0@(!_, (!_, !_)) -> run interval0 l r + Nothing -> pure r + + -- Looks into intervals with @l' >= l0@. + -- [----] + -- (i) *--------] overwrite if it's x + -- (ii) [-------]* delete anyways + -- (iii) *(------] overwrite if it's x, or + run (!l', (!r', !x')) l r + | l' > r = do + -- not adjacent: end. + pure r + -- (i) + | l' == r && x' == x = do + -- adjacent interval with the same value: merge into one. + onDel l' r' x' + IM.delete_ dim l' + pure r' + | l' == r = do + -- adjacent interval with different values: nothing to do. + pure r + -- (ii) + | r' <= r = do + -- inside the interval: delete and continue + onDel l' r' x' + IM.delete_ dim l' + res <- IM.lookupGT dim l' + case res of + Just rng -> run rng l r + Nothing -> pure r + -- (iii) + | x' == x = do + -- intersecting interval with the same value: merge into one. + onDel l' r' x' + 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 + + handleLeft l r = do + res <- IM.lookupLT dim l + case res of + Nothing -> pure (l, r) + Just (!l', (!r', !x')) + -- (i): adjacent interval + | 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) + | r' == l -> do + -- adjacent interval with different values: nothing to do. + pure (l, r) + -- (ii): not adjacent or intersecting + | r' < l -> do + pure (l, r) + -- (iii): intersecting + | 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') + | 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) + | 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) + +-- | Amortized \(O(\log n)\) Deletes an interval \([l, r)\) from the map. +-- +-- @since 1.1.0.0 +delete :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m () +delete itm l r = deleteM itm l r onAdd onDel + where + onAdd _ _ _ = pure () + onDel _ _ _ = pure () + +-- | Amortized \(O(\log n)\) Deletes an interval \([l, r)\) from the map. Tracks interval state +-- changes via @onAdd@ and @onDel@ hooks. +-- +-- @since 1.1.0.0 +deleteM :: + (PrimMonad m, VU.Unbox a) => + -- | The map + IntervalMap (PrimState m) a -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | @onAdd@ hook that take an interval \([l, r)\) with associated value \(v\) + (Int -> Int -> a -> m ()) -> + -- | @onDel@ hook that take an interval \([l, r)\) with associated value \(v\) + (Int -> Int -> a -> m ()) -> + m () +deleteM (IntervalMap dim) l0 r0 onAdd onDel + | l0 >= r0 = pure () + | otherwise = do + handleRight l0 r0 + handleLeft l0 r0 + where + handleRight l r = do + res <- IM.lookupGE dim l + case res of + Just interval0@(!_, (!_, !_)) -> run interval0 l r + Nothing -> pure () + + run (!l', (!r', !x')) l r + | l' >= r = do + -- not intersecting + pure () + | r' <= r = do + -- contained + onDel l' r' x' + IM.delete_ dim l' + res <- IM.lookupGT dim l' + case res of + Just rng -> run rng l r + Nothing -> pure () + | otherwise = do + -- intersecting + onDel l' r' x' + onAdd r r' x' + IM.delete_ dim l' + IM.insert dim r (r', x') + pure () + + handleLeft l r = do + res <- IM.lookupLT dim l + case res of + Nothing -> pure () + Just (!l', (!r', !x')) + | r' <= l -> do + -- not intersecting + pure () + | r' > r -> do + -- [l', r') contains [l, r) + 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') + | otherwise -> do + -- intersecting + onDel l' r' x' + onAdd l' l x' + -- IM.delete_ dim l' + 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 +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 + case res of + Just (!l', !r', !_) -> insert itm l' r' x + Nothing -> pure () + +-- | \(O(\log n)\). Shorthand for overwriting the value of an interval that contains \([l, r)\). +-- Tracks interval state changes via @onAdd@ and @onDel@ hooks. +-- +-- @since 1.1.0.0 +overwriteM :: + (PrimMonad m, Eq a, VU.Unbox a) => + -- | The map + IntervalMap (PrimState m) a -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(v\) + a -> + -- | @onAdd@ hook that take an interval \([l, r)\) with associated value \(v\) + (Int -> Int -> a -> m ()) -> + -- | @onDel@ hook that take an interval \([l, r)\) with associated value \(v\) + (Int -> Int -> a -> m ()) -> + m () +overwriteM itm l r x onAdd onDel = do + res <- lookup itm l r + case res of + Just (!l', !r', !_) -> insertM itm l' r' x onAdd onDel + Nothing -> pure () + +-- | \(O(n \log n)\) Enumerates the intervals and the associated values as \((l, (r, x))\) tuples, +-- where \([l, r)\) is the interval and \(x\) is the associated value. +-- +-- @since 1.1.0.0 +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/Math.hs b/src/AtCoder/Extra/Math.hs --- a/src/AtCoder/Extra/Math.hs +++ b/src/AtCoder/Extra/Math.hs @@ -1,34 +1,50 @@ -- | Extra math module. -- --- ==== Examples --- >>> import AtCoder.Extra.Math qualified as M --- >>> import Data.Semigroup (Product(..), Sum(..)) --- >>> getProduct $ M.power (<>) 32 (Product 2) --- 4294967296 --- --- >>> getProduct $ M.stimes' 32 (Product 2) --- 4294967296 --- --- >>> getProduct $ M.mtimes' 32 (Product 2) --- 4294967296 --- -- @since 1.0.0.0 module AtCoder.Extra.Math - ( -- * Binary exponential + ( -- * Re-exports from the internal math module + isPrime32, + ACIM.invGcd, + ACIM.primitiveRoot, + + -- * Binary exponentiation + + -- | ==== __Examples__ + -- >>> import AtCoder.Extra.Math qualified as M + -- >>> import Data.Semigroup (Product(..), Sum(..)) + -- >>> getProduct $ M.power (<>) 32 (Product 2) + -- 4294967296 + -- + -- >>> getProduct $ M.stimes' 32 (Product 2) + -- 4294967296 + -- + -- >>> getProduct $ M.mtimes' 32 (Product 2) + -- 4294967296 power, stimes', mtimes', ) where +import AtCoder.Internal.Math qualified as ACIM import Data.Bits ((.>>.)) --- TODO: add `HasCallStack` and provide with `unsafePower`. +-- | \(O(k \log^3 n) (k = 3)\). Returns whether the given `Int` value is a prime number. +-- +-- ==== Constraints +-- - \(n < 4759123141 (2^{32} < 4759123141)\), otherwise the return value can lie +-- (see [Wikipedia](https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test#Testing_against_small_sets_of_bases)). +-- +-- +-- @since 1.1.0.0 +{-# INLINE isPrime32 #-} +isPrime32 :: Int -> Bool +isPrime32 = ACIM.isPrime --- | Calculates \(s^n\) with custom multiplication operator using the binary exponentiation +-- | Calculates \(x^n\) with custom multiplication operator using the binary exponentiation -- technique. -- --- The internal implementation is taken from `Data.Semigroup.stimes`, but `power` uses strict +-- The internal implementation is taken from @Data.Semigroup.stimes@, but `power` uses strict -- evaluation and is often much faster. -- -- ==== Complexity @@ -53,7 +69,7 @@ | n == 1 = x `op` z | otherwise = g (x `op` x) (n .>>. 1) (x `op` z) --- | Strict `Data.Semigroup.stimes`. +-- | Strict variant of @Data.Semigroup.stimes@. -- -- ==== Complexity -- - \(O(\log n)\) @@ -66,7 +82,7 @@ stimes' :: (Semigroup a) => Int -> a -> a stimes' = power (<>) --- | Strict `Data.Monoid.mtimes`. +-- | Strict variant of @Data.Monoid.mtimes@. -- -- ==== Complexity -- - \(O(\log n)\) diff --git a/src/AtCoder/Extra/Monoid.hs b/src/AtCoder/Extra/Monoid.hs --- a/src/AtCoder/Extra/Monoid.hs +++ b/src/AtCoder/Extra/Monoid.hs @@ -1,31 +1,42 @@ {-# LANGUAGE TypeFamilies #-} --- | Extra module of pre-defined `SegAct` instances. --- --- Be warned that they're not 100% guaranteed to be correct. +-- | Extra module of pre-defined `SegAct` instances and helpful monoids. -- -- @since 1.0.0.0 module AtCoder.Extra.Monoid - ( -- * SegAct (re-export) + ( -- * Re-exports + + -- | It's mainly a list. It is recommended to use specific submodules. + + -- ** SegAct SegAct (..), - -- * Affine1 + -- ** Affine1 Affine1 (..), Affine1Repr, - -- * Range add + -- ** Mat2x2 + Mat2x2 (..), + Mat2x2Repr, + V2 (..), + V2Repr, + + -- ** Range add RangeAdd (..), - RangeAddId (..), - -- * Range set + -- ** Range set RangeSet (..), - RangeSetId (..), + RangeSetRepr, + + -- ** Rolling hash + RollingHash, ) where import AtCoder.Extra.Monoid.Affine1 (Affine1 (..), Affine1Repr) +import AtCoder.Extra.Monoid.Mat2x2 (Mat2x2 (..), Mat2x2Repr) import AtCoder.Extra.Monoid.RangeAdd (RangeAdd (..)) -import AtCoder.Extra.Monoid.RangeAddId (RangeAddId (..)) -import AtCoder.Extra.Monoid.RangeSet (RangeSet (..)) -import AtCoder.Extra.Monoid.RangeSetId (RangeSetId (..)) +import AtCoder.Extra.Monoid.RangeSet (RangeSet (..), RangeSetRepr) +import AtCoder.Extra.Monoid.RollingHash (RollingHash) +import AtCoder.Extra.Monoid.V2 (V2 (..), V2Repr) import AtCoder.LazySegTree (SegAct (..)) diff --git a/src/AtCoder/Extra/Monoid/Affine1.hs b/src/AtCoder/Extra/Monoid/Affine1.hs --- a/src/AtCoder/Extra/Monoid/Affine1.hs +++ b/src/AtCoder/Extra/Monoid/Affine1.hs @@ -1,17 +1,23 @@ {-# LANGUAGE TypeFamilies #-} --- | Range add monoid action for \([l, r)\) intervals: \(f: x \rightarrow ax + b\). +-- | Monoid action \(f: x \rightarrow ax + b\). -- +-- - Use @Mat2x2@ if inverse operations are required, or if it's necessary to store the monoid +-- length in the acted monoid (@V2@). +-- -- @since 1.0.0.0 module AtCoder.Extra.Monoid.Affine1 ( -- * Affine1 Affine1 (..), Affine1Repr, - -- * Constructor + -- * Constructors new, + unAffine1, + ident, + zero, - -- * Action + -- * Actions act, ) where @@ -27,17 +33,16 @@ import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM --- Tuple is not the fastest representation, but it's easier to implement `Unbox`. - --- | Range add monoid action for \([l, r)\) intervals: \(f: x \rightarrow ax + b\). +-- | Monoid action \(f: x \rightarrow ax + b\). -- +-- - Use @Mat2x2@ if inverse operations are required, or if it's necessary to store the monoid +-- length in the acted monoid (@V2@). +-- -- ==== Composition and dual --- `Semigroup` for `Affine1` is implemented like function composition, and rightmost affine --- transformation is applied first: \((f_1 \circ f_2) v := f_1 (f_2(v))\). If you need 'foldr' --- of \([f_l, f_{l+1}, .., f_r)\) on a segment tree, be sure to wrap `Affine1` in --- `Data.Monoid.Dual`. +-- The affine transformation acts as a left monoid action: \(f_2 (f_1 v) = (f_2 \circ f_1) v\). To +-- apply the leftmost transformation first in a segment tree, wrap `Affine1` in @Data.Monoid.Dual@. -- --- ==== Example +-- ==== __Example__ -- >>> import AtCoder.Extra.Monoid (SegAct(..), Affine1(..)) -- >>> import AtCoder.LazySegTree qualified as LST -- >>> seg <- LST.build @_ @(Affine1 Int) @(Sum Int) $ VU.generate 3 Sum -- [0, 1, 2] @@ -62,21 +67,42 @@ -- @since 1.0.0.0 type Affine1Repr a = (a, a) --- | Creates `Affine1`. +-- | \(O(1)\) Creates a one-dimensional affine transformation: \(f: x \rightarrow a \times x + b\). -- -- @since 1.0.0.0 {-# INLINE new #-} new :: a -> a -> Affine1 a new !a !b = Affine1 (a, b) --- | Applies \(f: x \rightarrow a \times x + b\). +-- | \(O(1)\) Retrieves the two components of `Affine1`. -- +-- @since 1.1.0.0 +{-# INLINE unAffine1 #-} +unAffine1 :: Affine1 a -> Affine1Repr a +unAffine1 (Affine1 a) = a + +-- | \(O(1)\) Identity transformation. +-- +-- @since 1.1.0.0 +{-# INLINE ident #-} +ident :: (Num a) => Affine1 a +ident = Affine1 (1, 0) + +-- | \(O(1)\) Transformation to zero. +-- +-- @since 1.1.0.0 +{-# INLINE zero #-} +zero :: (Num a) => Affine1 a +zero = Affine1 (0, 0) + +-- | \(O(1)\) Applies the one-dimensional affine transformation \(f: x \rightarrow a \times x + b\). +-- -- @since 1.0.0.0 {-# INLINE act #-} act :: (Num a) => Affine1 a -> a -> a act (Affine1 (!a, !b)) x = a * x + b --- | Acts on @a@ with length in terms of `SegAct`. Works for `Sum a` only. +-- | \(O(1)\) Acts on @a@ with length in terms of `SegAct`. Works for `Sum a` only. -- -- @since 1.0.0.0 {-# INLINE actWithLength #-} diff --git a/src/AtCoder/Extra/Monoid/Mat2x2.hs b/src/AtCoder/Extra/Monoid/Mat2x2.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Monoid/Mat2x2.hs @@ -0,0 +1,202 @@ +{-# LANGUAGE TypeFamilies #-} + +-- | Monoid action \(f: x \rightarrow ax + b\). Less efficient than @Affine1@, but compatible with +-- inverse opereations. +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Monoid.Mat2x2 + ( -- * Mat2x2 + Mat2x2 (..), + Mat2x2Repr, + + -- * Constructors + new, + unMat2x2, + ident, + zero, + + -- * Actions + act, + + -- * Operators + map, + det, + inv, + ) +where + +import AtCoder.Extra.Math qualified as ACEM +import AtCoder.Extra.Monoid.V2 (V2 (..)) +import AtCoder.LazySegTree (SegAct (..)) +import Data.Semigroup (Dual (..), Semigroup (..)) +import Data.Vector.Generic qualified as VG +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import GHC.Stack (HasCallStack) +import Prelude hiding (map) + +-- | Monoid action \(f: x \rightarrow ax + b\). Less efficient than @Affine1@, but compatible with +-- inverse opereations. +-- +-- ==== Composition and dual +-- The affine transformation acts as a left monoid action: \(f_2 (f_1 v) = (f_2 \circ f_1) v\). To +-- apply the leftmost transformation first in a segment tree, wrap `Mat2x2` in @Data.Monoid.Dual@. +-- +-- ==== __Example__ +-- >>> import AtCoder.Extra.Monoid.Mat2x2 qualified as Mat2x2 +-- >>> import AtCoder.Extra.Monoid.V2 qualified as V2 +-- >>> import AtCoder.Extra.Monoid (SegAct(..), Mat2x2(..), V2(..)) +-- >>> import AtCoder.LazySegTree qualified as LST +-- >>> seg <- LST.build @_ @(Mat2x2 Int) @(V2 Int) $ VU.generate 3 V2.new -- [0, 1, 2] +-- >>> LST.applyIn seg 0 3 $ Mat2x2.new 2 1 -- [1, 3, 5] +-- >>> V2.unV2 <$> LST.allProd seg +-- 9 +-- +-- @since 1.1.0.0 +newtype Mat2x2 a = Mat2x2 (Mat2x2Repr a) + deriving newtype + ( -- | @since 1.1.0.0 + Eq, + -- | @since 1.1.0.0 + Ord, + -- | @since 1.1.0.0 + Show + ) + +-- | `Mat2x2` internal representation. Tuples are not the fastest representation, but it's easier +-- to implement `Data.Vector.Unboxed.Unbox`. +-- +-- @since 1.1.0.0 +type Mat2x2Repr a = (a, a, a, a) + +-- | \(O(1)\) Creates a one-dimensional affine transformation: \(f: x \rightarrow a \times x + b\). +-- +-- @since 1.1.0.0 +{-# INLINE new #-} +new :: (Num a) => a -> a -> Mat2x2 a +new !a !b = Mat2x2 (a, b, 0, 1) + +-- | \(O(1)\) Retrieves the four components of `Mat2x2`. +-- +-- @since 1.1.0.0 +{-# INLINE unMat2x2 #-} +unMat2x2 :: Mat2x2 a -> Mat2x2Repr a +unMat2x2 (Mat2x2 a) = a + +-- | \(O(1)\) Transformation to zero. +-- +-- @since 1.1.0.0 +{-# INLINE zero #-} +zero :: (Num a) => Mat2x2 a +zero = Mat2x2 (0, 0, 0, 0) + +-- | \(O(1)\) Identity transformation. +-- +-- @since 1.1.0.0 +{-# INLINE ident #-} +ident :: (Num a) => Mat2x2 a +ident = Mat2x2 (1, 0, 0, 1) + +-- | \(O(1)\) Multiplies `Mat2x2` to `V2`. +{-# INLINE mulMV #-} +mulMV :: (Num a) => Mat2x2 a -> V2 a -> V2 a +mulMV (Mat2x2 (!a11, !a12, !a21, !a22)) (V2 (!x1, !x2)) = V2 (a', b') + where + !a' = a11 * x1 + a12 * x2 + !b' = a21 * x1 + a22 * x2 + +-- | \(O(1)\) Multiplies `Mat2x2` to `Mat2x2`. +{-# INLINE mulMM #-} +mulMM :: (Num a) => Mat2x2 a -> Mat2x2 a -> Mat2x2 a +mulMM (Mat2x2 (!a11, !a12, !a21, !a22)) (Mat2x2 (!b11, !b12, !b21, !b22)) = Mat2x2 (c11, c12, c21, c22) + where + !c11 = a11 * b11 + a12 * b21 + !c12 = a11 * b12 + a12 * b22 + !c21 = a21 * b11 + a22 * b21 + !c22 = a21 * b12 + a22 * b22 + +-- | \(O(1)\) Multiplies `Mat2x2` to `V2`. +-- +-- @since 1.1.0.0 +{-# INLINE act #-} +act :: (Num a) => Mat2x2 a -> V2 a -> V2 a +act = mulMV + +-- | \(O(1)\) Maps the every component of `Mat2x2`. +-- +-- @since 1.1.0.0 +{-# INLINE map #-} +map :: (a -> b) -> Mat2x2 a -> Mat2x2 b +map f (Mat2x2 (!a11, !a12, !a21, !a22)) = Mat2x2 (a11', a12', a21', a22') + where + !a11' = f a11 + !a12' = f a12 + !a21' = f a21 + !a22' = f a22 + +-- | \(O(1)\) Returns the determinan of the matrix. +-- +-- @since 1.1.0.0 +{-# INLINE det #-} +det :: (Fractional e) => Mat2x2 e -> e +det (Mat2x2 (!a, !b, !c, !d)) = a * d - b * c + +-- | \(O(1)\) Returns the inverse matrix, based on `Fractional` instance (mainly for @ModInt@). +-- +-- ==== Constraints +-- - The determinant (`det`) of the matrix must be non-zero, otherwise an error is thrown. +-- +-- @since 1.1.0.0 +{-# INLINE inv #-} +inv :: (HasCallStack, Fractional e, Eq e) => Mat2x2 e -> Mat2x2 e +inv (Mat2x2 (!a, !b, !c, !d)) = Mat2x2 (a', b', c', d') + where + -- NOTE: zero division + -- !r = recip $ a * d - b * c + !r + | det_ == 0 = error "AtCoder.Extra.Mat2x2.inv: the determinant of the matrix must be non zero" + | otherwise = recip det_ + where + !det_ = a * d - b * c + !a' = r * d + !b' = r * (-b) + !c' = r * (-c) + !d' = r * a + +-- | @since 1.1.0.0 +instance (Num a) => Semigroup (Mat2x2 a) where + {-# INLINE (<>) #-} + (<>) = mulMM + {-# INLINE stimes #-} + stimes = ACEM.stimes' . fromIntegral + +-- | @since 1.1.0.0 +instance (Num a) => Monoid (Mat2x2 a) where + {-# INLINE mempty #-} + mempty = ident + +-- | @since 1.1.0.0 +instance (Num a) => SegAct (Mat2x2 a) (V2 a) where + {-# INLINE segAct #-} + segAct = mulMV + +-- | @since 1.1.0.0 +instance (Num a) => SegAct (Dual (Mat2x2 a)) (V2 a) where + {-# INLINE segAct #-} + segAct (Dual f) = mulMV f + +-- | @since 1.1.0.0 +newtype instance VU.MVector s (Mat2x2 a) = MV_Mat2x2 (VU.MVector s (Mat2x2Repr a)) + +-- | @since 1.1.0.0 +newtype instance VU.Vector (Mat2x2 a) = V_Mat2x2 (VU.Vector (Mat2x2Repr a)) + +-- | @since 1.1.0.0 +deriving instance (VU.Unbox a) => VGM.MVector VUM.MVector (Mat2x2 a) + +-- | @since 1.1.0.0 +deriving instance (VU.Unbox a) => VG.Vector VU.Vector (Mat2x2 a) + +-- | @since 1.1.0.0 +instance (VU.Unbox a) => VU.Unbox (Mat2x2 a) diff --git a/src/AtCoder/Extra/Monoid/RangeAdd.hs b/src/AtCoder/Extra/Monoid/RangeAdd.hs --- a/src/AtCoder/Extra/Monoid/RangeAdd.hs +++ b/src/AtCoder/Extra/Monoid/RangeAdd.hs @@ -1,6 +1,6 @@ {-# LANGUAGE TypeFamilies #-} --- | Range add monoid action for \([l, r)\) intervals. +-- | Monoid action \(f: x \rightarrow x + d\). -- -- @since 1.0.0.0 module AtCoder.Extra.Monoid.RangeAdd @@ -9,27 +9,28 @@ -- * Constructor new, + unRangeAdd, - -- * Action + -- * Actions act, ) where import AtCoder.LazySegTree (SegAct (..)) -import Data.Semigroup (Sum (..)) +import Data.Semigroup (stimes, Sum (..), Max(..), Min(..)) 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 --- | Range add monoid action. +-- | Monoid action \(f: x \rightarrow x + d\). -- --- ==== Example +-- ==== __Example__ -- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeAdd(..)) -- >>> import AtCoder.LazySegTree qualified as LST -- >>> import Data.Semigroup (Max(..)) --- >>> seg <- LST.build @_ @(RangeAdd Int) @(Sum Int) $ VU.generate 3 Sum -- [0, 1, 2] --- >>> LST.applyIn seg 0 3 $ RangeAdd 5 -- [5, 6, 7] +-- >>> seg <- LST.build @_ @(RangeAdd (Sum Int)) @(Sum Int) $ VU.generate 3 Sum -- [0, 1, 2] +-- >>> LST.applyIn seg 0 3 $ RangeAdd (Sum 5) -- [5, 6, 7] -- >>> getSum <$> LST.prod seg 0 3 -- 18 -- @@ -51,36 +52,51 @@ new :: a -> RangeAdd a new = RangeAdd --- | Applies one-length range add: \(f: x \rightarrow d + x\). +-- | \(O(1)\) Retrieves the internal value of `RangeAdd`. -- +-- @since 1.1.0.0 +{-# INLINE unRangeAdd #-} +unRangeAdd :: RangeAdd a -> a +unRangeAdd (RangeAdd a) = a + +-- | \(O(1)\) Applies one-length range add: \(f: x \rightarrow d + x\). +-- -- @since 1.0.0.0 {-# INLINE act #-} -act :: (Num a) => RangeAdd a -> a -> a -act (RangeAdd dx) x = dx + x +act :: (Semigroup a) => RangeAdd a -> a -> a +act (RangeAdd dx) x = dx <> x --- | Acts on @a@ with length in terms of `SegAct`. +-- | \(O(1)\) Acts on @a@ with length in terms of `SegAct`. -- -- @since 1.0.0.0 {-# INLINE actWithLength #-} -actWithLength :: (Num a) => Int -> RangeAdd a -> a -> a -actWithLength len (RangeAdd f) x = fromIntegral len * f + x +actWithLength :: (Semigroup a) => Int -> RangeAdd a -> a -> a +actWithLength len (RangeAdd f) x = stimes len f <> x -- | @since 1.0.0.0 -instance (Num a) => Semigroup (RangeAdd a) where +instance (Semigroup a) => Semigroup (RangeAdd a) where {-# INLINE (<>) #-} - (RangeAdd a) <> (RangeAdd b) = RangeAdd $! a + b + (RangeAdd a) <> (RangeAdd b) = RangeAdd $! a <> b --- | @since 1.0.0.0 -instance (Num a) => Monoid (RangeAdd a) where +-- | @since 1.1.0.0 +instance (Monoid a) => Monoid (RangeAdd a) where {-# INLINE mempty #-} - mempty = RangeAdd 0 + mempty = RangeAdd mempty --- | @since 1.0.0.0 -instance (Num a) => SegAct (RangeAdd a) (Sum a) where +-- | @since 1.1.0.0 +instance (Monoid (Sum a)) => SegAct (RangeAdd (Sum a)) (Sum a) where {-# INLINE segActWithLength #-} - segActWithLength len a (Sum x) = Sum $! actWithLength len a x + segActWithLength len f x = actWithLength len f x --- not works as SegAct for Product, Min, and Max. +-- | @since 1.1.0.0 +instance (Monoid (Max a)) => SegAct (RangeAdd (Max a)) (Max a) where + {-# INLINE segActWithLength #-} + segActWithLength len f x = actWithLength len f x + +-- | @since 1.1.0.0 +instance (Monoid (Min a)) => SegAct (RangeAdd (Min a)) (Min a) where + {-# INLINE segActWithLength #-} + segActWithLength len f x = actWithLength len f x -- | @since 1.0.0.0 newtype instance VU.MVector s (RangeAdd a) = MV_RangeAdd (VU.MVector s a) diff --git a/src/AtCoder/Extra/Monoid/RangeAddId.hs b/src/AtCoder/Extra/Monoid/RangeAddId.hs deleted file mode 100644 --- a/src/AtCoder/Extra/Monoid/RangeAddId.hs +++ /dev/null @@ -1,95 +0,0 @@ -{-# LANGUAGE TypeFamilies #-} - --- | Range add monoid action for \([l, r)\) intervals. Works on ideomponent monoids such as `Max` --- or `Min` only. --- --- @since 1.0.0.0 -module AtCoder.Extra.Monoid.RangeAddId - ( -- * RangeAddId - RangeAddId (..), - -- * Constructor - new, - -- * Action - act, - ) -where - -import AtCoder.LazySegTree (SegAct (..)) -import Data.Semigroup (Max (..), Min (..)) -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 - --- | Range add monoid action. --- --- ==== Example --- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeAddId(..)) --- >>> import AtCoder.LazySegTree qualified as LST --- >>> import Data.Semigroup (Max(..)) --- >>> seg <- LST.build @_ @(RangeAddId Int) @(Max Int) $ VU.generate 3 Max -- [0, 1, 2] --- >>> LST.applyIn seg 0 3 $ RangeAddId 5 -- [5, 6, 7] --- >>> getMax <$> LST.prod seg 0 3 --- 7 --- --- @since 1.0.0.0 -newtype RangeAddId a = RangeAddId a - deriving newtype - ( -- | @since 1.0.0.0 - Eq, - -- | @since 1.0.0.0 - Ord, - -- | @since 1.0.0.0 - Show - ) - --- | Creates `RangeAddId`. --- --- @since 1.0.0.0 -{-# INLINE new #-} -new :: a -> RangeAddId a -new = RangeAddId - --- | Applies one-length range add: \(f: x \rightarrow d + x\). --- --- @since 1.0.0.0 -{-# INLINE act #-} -act :: (Num a) => RangeAddId a -> a -> a -act (RangeAddId f) x = f + x - --- | @since 1.0.0.0 -instance (Num a) => Semigroup (RangeAddId a) where - {-# INLINE (<>) #-} - (RangeAddId a) <> (RangeAddId b) = RangeAddId $! a + b - --- | @since 1.0.0.0 -instance (Num a) => Monoid (RangeAddId a) where - {-# INLINE mempty #-} - mempty = RangeAddId 0 - --- | @since 1.0.0.0 -instance (Num a) => SegAct (RangeAddId a) (Max a) where - {-# INLINE segAct #-} - segAct f (Max x) = Max $! act f x - --- | @since 1.0.0.0 -instance (Num a) => SegAct (RangeAddId a) (Min a) where - {-# INLINE segAct #-} - segAct f (Min x) = Min $! act f x - --- not works as SegAct for Sum and Product. - --- | @since 1.0.0.0 -newtype instance VU.MVector s (RangeAddId a) = MV_RangeAddId (VU.MVector s a) - --- | @since 1.0.0.0 -newtype instance VU.Vector (RangeAddId a) = V_RangeAddId (VU.Vector a) - --- | @since 1.0.0.0 -deriving instance (VU.Unbox a) => VGM.MVector VUM.MVector (RangeAddId a) - --- | @since 1.0.0.0 -deriving instance (VU.Unbox a) => VG.Vector VU.Vector (RangeAddId a) - --- | @since 1.0.0.0 -instance (VU.Unbox a) => VU.Unbox (RangeAddId a) diff --git a/src/AtCoder/Extra/Monoid/RangeSet.hs b/src/AtCoder/Extra/Monoid/RangeSet.hs --- a/src/AtCoder/Extra/Monoid/RangeSet.hs +++ b/src/AtCoder/Extra/Monoid/RangeSet.hs @@ -1,21 +1,22 @@ {-# LANGUAGE TypeFamilies #-} --- | Range set monoid action for \([l, r)\) intervals. +-- | Monoid action \(f: x \rightarrow a\). -- -- @since 1.0.0.0 module AtCoder.Extra.Monoid.RangeSet ( -- * RangeSet RangeSet (..), + RangeSetRepr, - -- * Constructor + -- * Constructors new, + unRangeSet, - -- * Action + -- * Actions act, ) where -import AtCoder.Extra.Math qualified as ACEM import AtCoder.LazySegTree (SegAct (..)) import Data.Bit (Bit (..)) import Data.Semigroup (stimes) @@ -24,9 +25,9 @@ import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM --- | Range set monoid action. +-- | Monoid action \(f: x \rightarrow a\). -- --- ==== Example +-- ==== __Example__ -- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeSet(..)) -- >>> import AtCoder.LazySegTree qualified as LST -- >>> import Data.Bit (Bit (..)) @@ -51,18 +52,25 @@ -- Tuples are not the fastest representation, but it's easier to implement -- `Data.Vector.Unboxed.Unbox`. -- --- @since 1.0.0.0 +-- @since 1.1.0.0 type RangeSetRepr a = (Bit, a) --- | Creates a new `RangeSet` action. +-- | \(O(1)\) Creates a new `RangeSet` action. -- -- @since 1.0.0.0 {-# INLINE new #-} new :: a -> RangeSet a new = RangeSet . (Bit True,) --- | Applies one-length range set: \(f: x \rightarrow y\). +-- | \(O(1)\) Retrieves the internal representation of `RangeSet`. -- +-- @since 1.1.0.0 +{-# INLINE unRangeSet #-} +unRangeSet :: RangeSet a -> RangeSetRepr a +unRangeSet (RangeSet a) = a + +-- | \(O(1)\) Applies one-length range set: \(f: x \rightarrow y\). +-- -- @since 1.0.0.0 {-# INLINE act #-} act :: RangeSet a -> a -> a @@ -74,7 +82,7 @@ -- @since 1.0.0.0 {-# INLINE actWithLength #-} actWithLength :: (Semigroup a) => Int -> RangeSet a -> a -> a -actWithLength len (RangeSet (Bit True, !f)) _ = ACEM.power (<>) len f +actWithLength len (RangeSet (Bit True, !f)) _ = stimes len f actWithLength _ (RangeSet (Bit False, !_)) x = x -- | @since 1.0.0.0 diff --git a/src/AtCoder/Extra/Monoid/RangeSetId.hs b/src/AtCoder/Extra/Monoid/RangeSetId.hs deleted file mode 100644 --- a/src/AtCoder/Extra/Monoid/RangeSetId.hs +++ /dev/null @@ -1,124 +0,0 @@ -{-# LANGUAGE TypeFamilies #-} - --- | Range set monoid action for \([l, r)\) intervals. Works on ideomponent monoids such as `Max` --- or `Min` only. --- --- --- @since 1.0.0.0 -module AtCoder.Extra.Monoid.RangeSetId - ( -- * RangeSetId - RangeSetId (..), - - -- * Constructor - new, - - -- * Action - act, - ) -where - -import AtCoder.LazySegTree (SegAct (..)) -import Data.Bit (Bit (..)) -import Data.Semigroup (Max (..), Min (..), stimes) -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 - --- | Range set monoid action. --- --- ==== Example --- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeSetId(..)) --- >>> import AtCoder.LazySegTree qualified as LST --- >>> import Data.Bit (Bit (..)) --- >>> import Data.Semigroup (Max(..)) --- >>> seg <- LST.build @_ @(RangeSetId (Max Int)) @(Max Int) $ VU.generate 3 (Max . (+ 10)) -- [10, 11, 12] --- >>> LST.applyIn seg 0 2 $ RangeSetId (Bit True, Max 5) -- [5, 5, 12] --- >>> getMax <$> LST.prod seg 0 3 --- 12 --- --- @since 1.0.0.0 -newtype RangeSetId a = RangeSetId (RangeSetIdRepr a) - deriving newtype - ( -- | @since 1.0.0.0 - Eq, - -- | @since 1.0.0.0 - Ord, - -- | @since 1.0.0.0 - Show - ) - --- | `RangeSetId` internal representation. The first value represents if it is an identity action. --- Tuples are not the fastest representation, but it's easier to implement --- `Data.Vector.Unboxed.Unbox`. --- --- @since 1.0.0.0 -type RangeSetIdRepr a = (Bit, a) - --- | Creates a new `RangeSet` action. --- --- @since 1.0.0.0 -{-# INLINE new #-} -new :: a -> RangeSetId a -new = RangeSetId . (Bit True,) - --- | Applies one-length range set: \(f: x \rightarrow y\). --- --- @since 1.0.0.0 -{-# INLINE act #-} -act :: RangeSetId a -> a -> a -act (RangeSetId (Bit True, !f)) _ = f -act (RangeSetId (Bit False, !_)) x = x - --- segActWithLength works for ideomponent monoids only. - --- | @since 1.0.0.0 -instance Semigroup (RangeSetId a) where - {-# INLINE (<>) #-} - RangeSetId (Bit False, !_) <> old = old - new_ <> _ = new_ - {-# INLINE stimes #-} - stimes _ x = x - --- The `Monoid` constraint is just for their default value. - --- | @since 1.0.0.0 -instance (Monoid a) => Monoid (RangeSetId a) where - {-# INLINE mempty #-} - mempty = RangeSetId (Bit False, mempty) - {-# INLINE mconcat #-} - -- find the first non-mempty - mconcat [] = mempty - mconcat (RangeSetId (Bit False, !_) : as) = mconcat as - mconcat (a : _) = a - --- The target is limited to ideomponent monoids. The `Monoid` constraint is just for their default --- value. - --- | @since 1.0.0.0 -instance (Ord a, Bounded a) => SegAct (RangeSetId (Max a)) (Max a) where - {-# INLINE segAct #-} - segAct = act - --- The target is limited to ideomponent monoids. The `Monoid` constraint is just for their default --- value. - --- | @since 1.0.0.0 -instance (Ord a, Bounded a) => SegAct (RangeSetId (Min a)) (Min a) where - {-# INLINE segAct #-} - segAct = act - --- | @since 1.0.0.0 -newtype instance VU.MVector s (RangeSetId a) = MV_RangeSetId (VU.MVector s (RangeSetIdRepr a)) - --- | @since 1.0.0.0 -newtype instance VU.Vector (RangeSetId a) = V_RangeSetId (VU.Vector (RangeSetIdRepr a)) - --- | @since 1.0.0.0 -deriving instance (VU.Unbox a) => VGM.MVector VUM.MVector (RangeSetId a) - --- | @since 1.0.0.0 -deriving instance (VU.Unbox a) => VG.Vector VU.Vector (RangeSetId a) - --- | @since 1.0.0.0 -instance (VU.Unbox a) => VU.Unbox (RangeSetId a) diff --git a/src/AtCoder/Extra/Monoid/RollingHash.hs b/src/AtCoder/Extra/Monoid/RollingHash.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Monoid/RollingHash.hs @@ -0,0 +1,121 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Rolling hash algorithm implemented as a monoid, typically stored in a segment tree. The type +-- parameters \(b\) and \(p\) represent the B-adic base and the modulus, respectively. +-- +-- Combining `RollingHash` with `SegTree` enables \(O(\log |s|)\) string slice creation and +-- \(O(1)\) slice comparison. +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Monoid.RollingHash + ( -- * Rolling hash + RollingHash (..), + + -- * Constructors + new, + unsafeNew, + ) +where + +import Data.Vector.Generic qualified as VG +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import GHC.Exts (proxy#) +import GHC.TypeNats (KnownNat, natVal') + +-- | Rolling hash algorithm implemented as a monoid, typically stored in a segment tree. The type +-- parameters \(b\) and \(p\) represent the B-adic base and the modulus, respectively. +-- +-- Combining `RollingHash` with `SegTree` enables \(O(\log |s|)\) string slice creation and +-- \(O(1)\) slice comparison. +-- +-- +-- ==== __Example__ +-- It's convenient to define a type alias of `RollingHash`: +-- +-- >>> import AtCoder.Extra.Monoid.RollingHash qualified as RH +-- >>> import AtCoder.SegTree qualified as ST +-- >>> import Data.Char (ord) +-- >>> import Data.Semigroup (Dual (..)) +-- >>> type RH = RH.RollingHash 100 998244353 +-- +-- Let's test whether "abcba" is a palindrome: +-- +-- >>> seg <- ST.build @_ @RH . VU.map (RH.unsafeNew . ord) $ VU.fromList "abcba" +-- >>> seg' <- ST.build @_ @(Dual RH) . VU.map (Dual . RH.unsafeNew . ord) $ VU.fromList "abcba" +-- >>> hash1 <- ST.prod seg 2 5 -- cba (left to right) +-- >>> Dual hash2 <- ST.prod seg' 0 3 -- abc (right to lett) +-- >>> hash1 == hash2 +-- True +-- +-- @since 1.1.0.0 +data RollingHash b p = RollingHash + { -- | The hash value. + hashRH :: {-# UNPACK #-} !Int, + -- | \(b^{\mathrm{length}} \bmod p\). + nextDigitRH :: {-# UNPACK #-} !Int + } + deriving + ( -- | @since 1.1.0.0 + Eq, + -- | @since 1.1.0.0 + Show + ) + +-- | \(O(1)\) Creates a one-length `RollingHash` from an integer. +-- +-- @since 1.1.0.0 +{-# INLINE new #-} +new :: forall b p. (KnownNat b, KnownNat p) => Int -> RollingHash b p +new h = RollingHash (h `mod` fromIntegral (natVal' (proxy# @p))) (fromIntegral (natVal' (proxy# @b))) + +-- | \(O(1)\) Creates a one-length `RollingHash` from an integer without taking the mod. +-- +-- @since 1.1.0.0 +{-# INLINE unsafeNew #-} +unsafeNew :: forall b p. (KnownNat b, KnownNat p) => Int -> RollingHash b p +unsafeNew h = RollingHash h (fromIntegral (natVal' (proxy# @b))) + +-- | @since 1.1.0.0 +instance (KnownNat b, KnownNat p) => Semigroup (RollingHash b p) where + -- \| \(O(1)\) + {-# INLINE (<>) #-} + (RollingHash !digit1 !hash1) <> (RollingHash !digit2 !hash2) = RollingHash digit' hash' + where + !p = fromIntegral $ natVal' (proxy# @p) + !digit' = digit1 * digit2 `mod` p + !hash' = (hash1 * digit2 + hash2) `mod` p + +-- | @since 1.1.0.0 +instance (KnownNat b, KnownNat p) => Monoid (RollingHash b p) where + {-# INLINE mempty #-} + mempty = RollingHash 1 0 + +type RHRepr = (Int, Int) + +-- | @since 1.1.0.0 +instance VU.IsoUnbox (RollingHash b p) RHRepr where + {-# INLINE toURepr #-} + toURepr (RollingHash a b) = (a, b) + {-# INLINE fromURepr #-} + fromURepr (!a, !b) = RollingHash a b + +-- | @since 1.1.0.0 +newtype instance VU.MVector s (RollingHash b p) = MV_RH (VUM.MVector s RHRepr) + +-- | @since 1.1.0.0 +newtype instance VU.Vector (RollingHash b p) = V_RH (VU.Vector RHRepr) + +-- | @since 1.1.0.0 +deriving via (RollingHash b p `VU.As` RHRepr) instance VGM.MVector VUM.MVector (RollingHash b p) + +-- | @since 1.1.0.0 +deriving via (RollingHash b p `VU.As` RHRepr) instance VG.Vector VU.Vector (RollingHash b p) + +-- | @since 1.1.0.0 +instance VU.Unbox (RollingHash b p) diff --git a/src/AtCoder/Extra/Monoid/V2.hs b/src/AtCoder/Extra/Monoid/V2.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Monoid/V2.hs @@ -0,0 +1,81 @@ +{-# LANGUAGE TypeFamilies #-} + +-- | A monoid acted on by `Mat2x2`, an affine transformation target. +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Monoid.V2 + ( -- * V2 + V2 (..), + V2Repr, + + -- * Constructor + new, + unV2, + ) +where + +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 + +-- | A monoid acted on by `Mat2x2`, an affine transformation target. +-- +-- @since 1.1.0.0 +newtype V2 a = V2 (V2Repr a) + deriving newtype + ( -- | @since 1.1.0.0 + Eq, + -- | @since 1.1.0.0 + Ord, + -- | @since 1.1.0.0 + Show + ) + +-- | `V2` internal representation. Tuples are not the fastest representation, but it's easier +-- to implement `Data.Vector.Unboxed.Unbox`. +-- +-- @since 1.1.0.0 +type V2Repr a = (a, a) + +-- | \(O(1)\) Creates `V2` of length \(1\). +-- +-- @since 1.1.0.0 +{-# INLINE new #-} +new :: (Num a) => a -> V2 a +new !a = V2 (a, 1) + +-- | \(O(1)\) Retrieves the value of `V2`, discarding the length information. +-- +-- @since 1.1.0.0 +{-# INLINE unV2 #-} +unV2 :: V2 a -> a +unV2 (V2 (!a, !_)) = a + +-- | @since 1.1.0.0 +instance (Num a) => Semigroup (V2 a) where + {-# INLINE (<>) #-} + (V2 (!a1, !a2)) <> (V2 (!b1, !b2)) = V2 (a', b') + where + !a' = a1 + b1 + !b' = a2 + b2 + +-- | @since 1.1.0.0 +instance (Num a) => Monoid (V2 a) where + {-# INLINE mempty #-} + mempty = V2 (0, 0) + +-- | @since 1.1.0.0 +newtype instance VU.MVector s (V2 a) = MV_V2 (VU.MVector s (V2Repr a)) + +-- | @since 1.1.0.0 +newtype instance VU.Vector (V2 a) = V_V2 (VU.Vector (V2Repr a)) + +-- | @since 1.1.0.0 +deriving instance (VU.Unbox a) => VGM.MVector VUM.MVector (V2 a) + +-- | @since 1.1.0.0 +deriving instance (VU.Unbox a) => VG.Vector VU.Vector (V2 a) + +-- | @since 1.1.0.0 +instance (VU.Unbox a) => VU.Unbox (V2 a) diff --git a/src/AtCoder/Extra/MultiSet.hs b/src/AtCoder/Extra/MultiSet.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/MultiSet.hs @@ -0,0 +1,280 @@ +{-# LANGUAGE RecordWildCards #-} + +-- | A fast, mutable multiset for `Int` keys backed by a @HashMap@. Most operations are performed +-- in \(O(1)\) time, but in average. +-- +-- ==== Invariant +-- The count for each key must be non-negative. An exception is thrown if this invariant is +-- violated. +-- +-- ==== Capacity limitation +-- The maximum number of distinct keys that can be inserted is fixed at \(n\), even if some keys are +-- deleted. This is due to the limitation of the internal @HashMap@. +-- +-- ==== __Example__ +-- Create a `MultiSet` with capacity \(4\): +-- +-- >>> import AtCoder.Extra.MultiSet qualified as MS +-- >>> ms <- MS.new 4 +-- +-- `inc` and `dec` are the primary API: +-- +-- >>> MS.inc ms 10 +-- >>> MS.inc ms 10 +-- >>> MS.lookup ms 10 +-- Just 2 +-- +-- >>> MS.dec ms 10 +-- >>> MS.lookup ms 10 +-- Just 1 +-- +-- Entries with zero count are considered to be non-existing: +-- +-- >>> MS.dec ms 10 +-- >>> MS.member ms 10 +-- False +-- +-- >>> MS.lookup ms 10 +-- Nothing +-- +-- >>> MS.size ms +-- 0 +-- +-- Creating a negative count results in an exception: +-- +-- >>> MS.inc ms 11 +-- >>> MS.sub ms 11 2 +-- *** Exception: AtCoder.Extra.Multiset.sub: the count of `11` is becoming a negative value: `-1` +-- ... +-- +-- Decrementing a non-existing key does nothing and does not throw an exception: +-- +-- >>> MS.dec ms 12 +-- +-- Misc: +-- +-- >>> MS.insert ms 12 112 +-- >>> MS.assocs ms +-- [(11,1),(12,112)] +-- +-- @since 1.1.0.0 +module AtCoder.Extra.MultiSet + ( -- * MultiSet + MultiSet, + + -- * Construtors + new, + + -- * Metadata + capacity, + size, + + -- * Lookups + lookup, + member, + notMember, + + -- * Modifications + inc, + dec, + add, + sub, + insert, + delete, + + -- * Conversions + + -- ** Safe conversions + keys, + elems, + assocs, + + -- ** Unsafe conversions + unsafeKeys, + unsafeElems, + unsafeAssocs, + ) +where + +import AtCoder.Extra.HashMap qualified as HM +import Control.Monad (when) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Data.Functor ((<&>)) +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 Prelude hiding (lookup) + +-- | A fast, mutable multiset for `Int` keys backed by a @HashMap@. +-- +-- @since 1.1.0.0 +data MultiSet s = MultiSet + { mapMS :: !(HM.HashMap s Int), + cntMS :: !(VUM.MVector s Int) + } + +-- | \(O(n)\) Creates a `MultiSet` with capacity \(n\). +-- +-- @since 1.1.0.0 +new :: (PrimMonad m) => Int -> m (MultiSet (PrimState m)) +new n = do + mapMS <- HM.new n + cntMS <- VUM.replicate 1 0 + pure $ MultiSet {..} + +-- | \(O(1)\) Returns the maximum number of distinct keys that can be inserted into the internal +-- hash map. +-- +-- @since 1.1.0.0 +capacity :: MultiSet s -> Int +capacity = HM.capacity . mapMS + +-- | \(O(1)\) Returns the number of distinct keys with positive counts. +-- +-- @since 1.1.0.0 +size :: (PrimMonad m) => MultiSet (PrimState m) -> m Int +size MultiSet {..} = do + VGM.unsafeRead cntMS 0 + +-- | \(O(1)\) Looks up the count for a key. +-- +-- @since 1.1.0.0 +lookup :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m (Maybe Int) +lookup MultiSet {..} k = do + HM.lookup mapMS k <&> \case + Just i | i > 0 -> Just i + _ -> Nothing + +-- | \(O(1)\) Tests whether \(k\) is in the set. +-- +-- @since 1.1.0.0 +member :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m Bool +member MultiSet {..} k = do + HM.lookup mapMS k <&> \case + Just i -> i > 0 + _ -> False + +-- | \(O(1)\) Tests whether \(k\) is not in the set. +-- +-- @since 1.1.0.0 +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 +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 +dec :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m () +dec ms k = sub ms k 1 + +-- | \(O(1)\) Increments the count of a key \(k\) by \(c\). If the key does not exist in the set, +-- the \((k, c)\) pair is inserted. If \(v\) is negative, it falls back to `sub`. +-- +-- @since 1.1.0.0 +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) + EQ -> pure () + GT -> do + HM.lookup mapMS k >>= \case + Just n -> do + HM.insert mapMS k $ n + v + when (n <= 0) $ do + VGM.unsafeModify cntMS (+ 1) 0 + Nothing -> do + HM.insert mapMS k v + VGM.unsafeModify cntMS (+ 1) 0 + +-- | \(O(1)\) Decrements the count of a key \(k\) by \(c\). If \(c\) is negative, it falls back to +-- `add`. +-- +-- @since 1.1.0.0 +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) + EQ -> pure () + GT -> do + HM.lookup mapMS k >>= \case + Just 0 -> pure () -- ignored + Just n -> case compare n v of + GT -> do + HM.insert mapMS k (n - v) + EQ -> do + HM.insert mapMS k 0 + VGM.unsafeModify cntMS (subtract 1) 0 + LT -> error $ "AtCoder.Extra.Multiset.sub: the count of `" ++ show k ++ "` is becoming a negative value: `" ++ show (n - v) ++ "`" + _ -> pure () + +-- | \(O(1)\) Inserts a key-count pair into the set. `MultiSet` is actually a count map. +-- +-- @since 1.1.0.0 +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 ++ "`" + | otherwise = do + HM.lookup mapMS k >>= \case + Just n | n > 0 -> do + HM.insert mapMS k v + _ -> do + HM.insert mapMS k v + VGM.unsafeModify cntMS (+ 1) 0 + +-- | \(O(1)\) Deletes a key. Note that it does not undo its insertion and does not increase the +-- number of distinct keys that can be inserted into the internal hash map. +-- +-- @since 1.1.0.0 +delete :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m () +delete MultiSet {..} k = do + HM.lookup mapMS k >>= \case + Just i | i > 0 -> do + HM.insert mapMS k 0 + VGM.unsafeModify cntMS (subtract 1) 0 + _ -> pure () + +-- | \(O(n)\) Enumerates the keys in the set. +-- +-- @since 1.1.0.0 +{-# INLINE keys #-} +keys :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int) +keys ms = VU.force <$> unsafeKeys ms + +-- | \(O(n)\) Enumerates the counts in the set. +-- +-- @since 1.1.0.0 +{-# INLINE elems #-} +elems :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int) +elems ms = VU.force <$> unsafeElems ms + +-- | \(O(n)\) Enumerates the key-count pairs in the set. +-- +-- @since 1.1.0.0 +{-# INLINE assocs #-} +assocs :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector (Int, Int)) +assocs ms = VU.force <$> unsafeAssocs ms + +-- | \(O(n)\) Enumerates the keys in the set. +-- +-- @since 1.1.0.0 +{-# INLINE unsafeKeys #-} +unsafeKeys :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int) +unsafeKeys = (VU.mapMaybe (\(!k, !n) -> if n > 0 then Just k else Nothing) <$>) . HM.unsafeAssocs . mapMS + +-- | \(O(n)\) Enumerates the counts in the set. +-- +-- @since 1.1.0.0 +{-# INLINE unsafeElems #-} +unsafeElems :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int) +unsafeElems = (VU.filter (> 0) <$>) . HM.unsafeElems . mapMS + +-- | \(O(n)\) Enumerates the key-count pairs in the set. +-- +-- @since 1.1.0.0 +{-# INLINE unsafeAssocs #-} +unsafeAssocs :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector (Int, Int)) +unsafeAssocs = (VU.filter (\(!_, !n) -> n > 0) <$>) . HM.unsafeAssocs . mapMS diff --git a/src/AtCoder/Extra/Pdsu.hs b/src/AtCoder/Extra/Pdsu.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Pdsu.hs @@ -0,0 +1,306 @@ +{-# LANGUAGE RecordWildCards #-} + +-- original implementation: +-- + +-- | A potentialized disjoint set union on a [group](https://en.wikipedia.org/wiki/Group_(mathematics\)) +-- under a differential constraint system. Each vertex \(v\) is assigned a potential value \(p(v)\), +-- where representatives (`leader`) of each group have a potential of `mempty`, and other vertices have +-- potentials relative to their representative. +-- +-- The group type is represented as a `Monoid` with a inverse operator, passed on `new`. This +-- approach avoids defining a separate typeclass for groups. +-- +-- ==== Invariant +-- New monoids always come from the left: @new <> old@. The order is important for non-commutative +-- monoid implementations. +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Pdsu + ( -- * Pdsu + Pdsu (nPdsu), + + -- * Constructors + new, + + -- * Inspection + leader, + pot, + diff, + unsafeDiff, + same, + canMerge, + + -- * Merging + merge, + merge_, + + -- * Group information + size, + groups, + + -- * Reset + clear, + ) +where + +import AtCoder.Internal.Assert qualified as ACIA +import Control.Monad +import Control.Monad.Primitive (PrimMonad, PrimState) +import Data.Vector qualified as V +import Data.Vector.Generic qualified as VG +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import GHC.Stack (HasCallStack) + +-- | A potentialized disjoint set union on a [group](https://en.wikipedia.org/wiki/Group_(mathematics\)) +-- under a differential constraint system. Each vertex \(v\) is assigned a potential value \(p(v)\), +-- +-- ==== __Example__ +-- Create a `Pdsu` with four vertices with potential type @Sum Int@. Use `negate` as the inverse +-- operator: +-- +-- >>> import AtCoder.Extra.Pdsu qualified as Pdsu +-- >>> import Data.Semigroup (Sum (..)) +-- >>> dsu <- Pdsu.new @_ @(Sum Int) 4 negate +-- +-- The API is similar to @Dsu@, but with differential potential values: +-- +-- >>> Pdsu.merge dsu 1 0 (Sum 1) -- p(1) - p(0) := Sum 1 +-- True +-- +-- >>> Pdsu.merge_ dsu 2 0 (Sum 2) -- p(2) - p(0) := Sum 2 +-- >>> Pdsu.leader dsu 0 +-- 0 +-- +-- Potential values can be retrieved with `pot`: +-- +-- >>> Pdsu.pot dsu 0 +-- Sum {getSum = 0} +-- +-- >>> Pdsu.pot dsu 1 +-- Sum {getSum = 1} +-- +-- >>> Pdsu.pot dsu 2 +-- Sum {getSum = 2} +-- +-- Difference of potentials in the same group can be retrieved with `diff`: +-- +-- >>> Pdsu.diff dsu 2 1 +-- Just (Sum {getSum = 1}) +-- +-- >>> Pdsu.diff dsu 2 3 +-- Nothing +-- +-- Retrieve group information with `groups` +-- +-- >>> Pdsu.groups dsu +-- [[2,1,0],[3]] +-- +-- @since 1.1.0.0 +data Pdsu s a = Pdsu + { -- | The number of vertices. + nPdsu :: {-# UNPACK #-} !Int, + -- | Parent: non-positive, size: positive + parentOrSizePdsu :: !(VUM.MVector s Int), + -- | Diffierencial potential of each vertex. + potentialPdsu :: !(VUM.MVector s a), + invertPdsu :: !(a -> a) + } + +-- | \(O(n)\) Creates a new DSU under a differential constraint system. +-- +-- @since 1.1.0.0 +{-# INLINE new #-} +new :: + forall m a. + (PrimMonad m, Monoid a, VU.Unbox a) => + -- | The number of vertices + Int -> + -- | The inverse operator of the monoid + (a -> a) -> + -- | A DSU + m (Pdsu (PrimState m) a) +new n f = Pdsu n <$> VUM.replicate n (-1 {- size 1 -}) <*> VUM.replicate n (mempty :: a) <*> pure f + +-- | \(O(\alpha(n))\) Returns the representative of the connected component that contains the +-- vertex. +-- +-- @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 + where + !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.leader" v0 nPdsu + inner v = do + p <- VGM.read parentOrSizePdsu v + if {- size? -} p < 0 + then pure v + else do + -- NOTE(perf): Path compression. + -- Handle the nodes closer to the root first and move them onto just under the root + !r <- inner p + when (p /= r) $ do + !pp <- VGM.read potentialPdsu p + -- Move `v` to just under the root: + VGM.write parentOrSizePdsu v {- root -} r + -- INVARIANT: new coming monoids always come from the left. And we're performing + -- reverse folding. + VGM.modify potentialPdsu (<> pp) v + pure r + +-- | \(O(\alpha(n))\) Returns \(p(v)\), the potential value of vertex \(v\) relative to the +-- reprensetative of its group. +-- +-- @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 + -- Perform path compression + _ <- leader dsu v1 + VGM.read potentialPdsu v1 + where + !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.pot" v1 nPdsu + +-- | \(O(\alpha(n))\) Returns whether the vertices \(a\) and \(b\) are in the same connected +-- component. +-- +-- @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 + +-- TODO: call it unsafeDiff + +-- | \(O(\alpha(n))\) Returns the potential of \(v_1\) relative to \(v_2\): \(p(v_1) \cdot p^{-1}(v_2)\) +-- if the two vertices belong to the same group. Returns `Nothing` when the two vertices are not +-- connected. +-- +-- @since 1.1.0.0 +{-# INLINE diff #-} +diff :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> m (Maybe a) +diff !dsu !v1 !v2 = do + b <- same dsu v1 v2 + if b + then Just <$> unsafeDiff dsu v1 v2 + else pure Nothing + +-- | \(O(\alpha(n))\) Returns the potential of \(v_1\) relative to \(v_2\): \(p(v_1) \cdot p^{-1}(v_2)\) +-- if the two vertices belong to the same group. Returns meaningless value if the two vertices are +-- not connected. +-- +-- @since 1.1.0.0 +{-# INLINE unsafeDiff #-} +unsafeDiff :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> m a +unsafeDiff !dsu !v1 !v2 = do + p1 <- pot dsu v1 + p2 <- pot dsu v2 + pure $ p1 <> invertPdsu dsu p2 + +-- | \(O(\alpha(n))\) Merges \(v_1\) to \(v_2\) with differential (relative) potential +-- \(\mathrm{dp}\): \(p(v1) := \mathrm{dp} \cdot p(v2)\). Returns `True` if they're newly merged. +-- +-- @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 + where + !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v10 nPdsu + !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v20 nPdsu + inner v1 v2 !dp = do + !r1 <- leader dsu v1 + !r2 <- leader dsu v2 + if r1 == r2 + then pure False + else do + -- NOTE(perf): Union by size (choose smaller one for root). + -- Another, more proper optimization would be union by rank (depth). + !size1 <- VGM.read potentialPdsu v1 + !size2 <- VGM.read potentialPdsu v2 + if size1 >= size2 + then do + -- Merge `r1` onto `r2` + + -- Update the size of `r1` + !negativeSize1 <- negate {- retrieve size -} <$> VGM.read parentOrSizePdsu r1 + !negativeSize2 <- negate {- retrieve size -} <$> VGM.read parentOrSizePdsu r2 + VGM.write parentOrSizePdsu r1 ({- size -} negativeSize1 + negativeSize2) + + -- p(v1) becomes p'(v1) under r2 after merge. p(r1) becomes p'(r1). + -- p'(v1) = dp <> p(v2) + -- p'(v1) = p(v1) <> 'p(r1) + -- Therefore, + -- p'(r1) = p^{-1}(v1) <> dp <> p(v2) + !p1 <- VGM.read potentialPdsu v1 + !p2 <- VGM.read potentialPdsu v2 + let !pr1' = invertPdsu p1 <> dp <> p2 + + -- Move `r1` to just under `r2`: + VGM.write parentOrSizePdsu r1 {- record new root -} r2 + VGM.write potentialPdsu r1 pr1' + + pure True + else do + inner v2 v1 $ invertPdsu dp + +-- | \(O(\alpha(n))\) `merge` with the return value discarded. +-- +-- @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 () +merge_ !dsu !v1 !v2 !dp = do + _ <- merge dsu v1 v2 dp + pure () + +-- | \(O(\alpha(n))\) Returns `True` if the two vertices belong to different groups or they belong +-- to the same group under the condition \(p(v_1) = dp \cdot p(v_2)\). It's just a convenient +-- helper function. +-- +-- @since 1.1.0.0 +{-# INLINE canMerge #-} +canMerge :: (HasCallStack, PrimMonad m, Semigroup a, Eq a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m Bool +canMerge !dsu !v1 !v2 !dp = do + b <- same dsu v1 v2 + if not b + then pure True + else do + !p1 <- VGM.read (potentialPdsu dsu) v1 + !p2 <- VGM.read (potentialPdsu dsu) v2 + pure $ p1 == dp <> p2 + +-- | \(O(\alpha(n))\) Returns the number of vertices belonging to the same group. +-- +-- @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 + +-- | \(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 + groupSize <- VUM.replicate nPdsu (0 :: Int) + leaders <- VU.generateM nPdsu $ \i -> do + li <- leader dsu i + VGM.modify groupSize (+ 1) li + pure li + result <- do + groupSize' <- VU.unsafeFreeze groupSize + V.mapM VUM.unsafeNew $ VU.convert groupSize' + VU.iforM_ leaders $ \i li -> do + i' <- subtract 1 <$> VGM.read groupSize li + VGM.write (result VG.! li) i' i + VGM.write groupSize li i' + V.filter (not . VU.null) <$> V.mapM VU.unsafeFreeze result + +-- | \(O(n)\) Clears the `Pdsu` to the initial state. +-- +-- @since 1.1.0.0 +{-# INLINE clear #-} +clear :: forall m a. (PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> m () +clear !dsu = do + VGM.set (potentialPdsu dsu) (mempty @a) + VGM.set (parentOrSizePdsu dsu) (-1 {- size -}) diff --git a/src/AtCoder/Extra/Semigroup/Matrix.hs b/src/AtCoder/Extra/Semigroup/Matrix.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Semigroup/Matrix.hs @@ -0,0 +1,264 @@ +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RecordWildCards #-} + +-- | A simple HxW matrix backed by a vector, mainly for binary exponention. +-- +-- The matrix is a left semigroup action: \(m_2 (m_1 v) = (m_2 \circ m_1) v\). +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Semigroup.Matrix + ( -- * Matrix + Matrix (..), + + -- * Constructors + new, + zero, + ident, + diag, + + -- * Mapping + map, + + -- * Multiplications + mulToCol, + mul, + mulMod, + mulMint, + + -- * Powers + pow, + powMod, + powMint, + ) +where + +import AtCoder.Extra.Math qualified as ACEM +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.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.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import GHC.Exts (proxy#) +import GHC.Stack (HasCallStack) +import GHC.TypeNats (KnownNat, natVal') +import Prelude hiding (map) + +-- | A simple HxW matrix backed by a vector, mainly for binary exponention. +-- +-- The matrix is a left semigroup action: \(m_2 (m_1 v) = (m_2 \circ m_1) v\). +-- +-- +-- @since 1.1.0.0 +data Matrix a = Matrix + { -- | @since 1.1.0.0 + hM :: {-# UNPACK #-} !Int, + -- | @since 1.1.0.0 + wM :: {-# UNPACK #-} !Int, + -- | @since 1.1.0.0 + vecM :: !(VU.Vector a) + } + deriving + ( -- | @since 1.1.0.0 + Show, + -- | @since 1.1.0.0 + Eq + ) + +-- | Type alias of a column vector. +-- +-- @since 1.1.0.0 +type Col a = VU.Vector a + +-- | \(O(hw)\) Creates an HxW matrix. +-- +-- @since 1.1.0.0 +{-# INLINE new #-} +new :: (HasCallStack, VU.Unbox a) => Int -> Int -> VU.Vector a -> Matrix a +new h w vec + | VU.length vec /= h * w = error "AtCoder.Extra.Matrix: size mismatch" + | otherwise = Matrix h w vec + +-- | \(O(n^2)\) Creates an NxN zero matrix. +-- +-- @since 1.1.0.0 +{-# INLINE zero #-} +zero :: (VU.Unbox a, Num a) => Int -> Matrix a +zero n = Matrix n n $ VU.replicate (n * n) 0 + +-- | \(O(n^2)\) Creates an NxN identity matrix. +-- +-- @since 1.1.0.0 +{-# INLINE ident #-} +ident :: (VU.Unbox a, Num a) => Int -> Matrix a +ident n = Matrix n n $ VU.create $ do + vec <- VUM.replicate (n * n) 0 + for_ [0 .. n - 1] $ \i -> do + VGM.write vec (i + n * i) 1 + pure vec + +-- | \(O(n^2)\) Creates an NxN diagonal matrix. +-- +-- @since 1.1.0.0 +{-# INLINE diag #-} +diag :: (VU.Unbox a, Num a) => Int -> VU.Vector a -> Matrix a +diag n xs = Matrix n n $ VU.create $ do + vec <- VUM.replicate (n * n) 0 + VU.iforM_ xs $ \i x -> do + VGM.write vec (i + n * i) x + pure vec + +-- | \(O(n^2)\) Maps the `Matrix`. +-- +-- @since 1.1.0.0 +{-# INLINE map #-} +map :: (VU.Unbox a, VU.Unbox b) => (a -> b) -> Matrix a -> Matrix b +map f Matrix {..} = Matrix hM wM $ VU.map f vecM + +-- | \(O(hw)\) Multiplies HxW matrix to a Hx1 column vector. +-- +-- @since 1.1.0.0 +{-# INLINE mulToCol #-} +mulToCol :: (Num a, VU.Unbox a) => Matrix a -> Col a -> Col a +mulToCol Matrix {..} !col = VU.convert $ V.map (VU.sum . VU.zipWith (*) col) rows + where + !n = VU.length col + !_ = ACIA.runtimeAssert (n == wM) "AtCoder.Extra.Matrix.mulToCol: size mismatch" + rows = V.unfoldrExactN hM (VU.splitAt wM) vecM + +-- | \(O(h_1 K w_2)\) Multiplies H1xK matrix to a KxW2 matrix. +-- +-- @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) + 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. +-- +-- @since 1.1.0.0 +{-# INLINE mulMod #-} +mulMod :: Int -> Matrix Int -> Matrix Int -> Matrix Int +mulMod !m !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) + 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 + !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMod: matrix size mismatch" + +-- | \(O(h_1 w_2 K)\) `mul` specialized to `M.ModInt`. +-- +-- @since 1.1.0.0 +{-# INLINE mulMint #-} +mulMint :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a) +mulMint = mulMintImpl bt + where + !bt = BT.new32 $ fromIntegral (natVal' (proxy# @a)) + +{-# 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 = + 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 + 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" + +-- | \(O(w n^3)\) Calculates \(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 = ident $ hM mat + | otherwise = ACEM.power mul k mat + where + !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMod: matrix size mismatch" + +-- | \(O(w n^3)\) Calculates \(M^k\), taking the mod. +-- +-- @since 1.1.0.0 +{-# INLINE powMod #-} +powMod :: Int -> Int -> Matrix Int -> Matrix Int +powMod m k mat + | k < 0 = error "AtCoder.Extra.Matrix.powMod: the exponential must be non-negative" + | k == 0 = ident $ hM mat + | otherwise = ACEM.power (mulMod m) k mat + 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`. +-- +-- @since 1.1.0.0 +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" + | k == 0 = ident $ hM mat + | otherwise = ACEM.power (mulMintImpl bt) k mat + where + !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMint: matrix size mismatch" + !bt = BT.new32 $ fromIntegral (natVal' (proxy# @m)) + +-- | @since 1.1.0.0 +instance (Num a, VU.Unbox a) => Semigroup (Matrix a) where + {-# INLINE (<>) #-} + (<>) = mul + {-# INLINE stimes #-} + stimes = ACEM.power (<>) . fromIntegral diff --git a/src/AtCoder/Extra/Semigroup/Permutation.hs b/src/AtCoder/Extra/Semigroup/Permutation.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Semigroup/Permutation.hs @@ -0,0 +1,114 @@ +-- | A permutation represented by a vector, mainly for binary exponentiation. +-- +-- The permutation is a left semigroup action: \(p_2 (p_1 x) = (p_2 \circ p_1) x\). +-- +-- ==== __Example__ +-- >>> import AtCoder.Extra.Semigroup.Permutation qualified as Permutation +-- >>> import Data.Semigroup (Semigroup (stimes)) +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let perm = Permutation.new $ VU.fromList [1, 2, 3, 0] +-- >>> Permutation.act perm 1 +-- 2 +-- +-- >>> Permutation.act (perm <> perm) 1 +-- 3 +-- +-- >>> Permutation.act (stimes 3 perm) 1 +-- 0 +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Semigroup.Permutation + ( -- * Permutation + Permutation (..), + + -- * Constructors + new, + unsafeNew, + ident, + zero, + + -- * Actions + act, + + -- * Metadata + length, + ) +where + +import AtCoder.Internal.Assert qualified as ACIA +import Data.Vector.Generic qualified as VG +import Data.Vector.Unboxed qualified as VU +import GHC.Stack (HasCallStack) +import Prelude hiding (length) + +-- | A permutation represented by a vector, mainly for binary exponentiation. +-- +-- The permutation is a left semigroup action: \(p_2 (p_1 x) = (p_2 \circ p_1) x\). +-- +-- @since 1.1.0.0 +newtype Permutation = Permutation + { unPermutation :: VU.Vector Int + } + deriving newtype + ( -- | @since 1.1.0.0 + Eq, + -- | @since 1.1.0.0 + Show + ) + +-- | \(O(1)\) Creates a `Permutation`, performing boundary check on input vector. +-- +-- @since 1.1.0.0 +{-# INLINE new #-} +new :: (HasCallStack) => VU.Vector Int -> Permutation +new xs = Permutation xs + where + n = VU.length xs + !_ = VU.foldl' (\() i -> let !_ = ACIA.runtimeAssert (-1 <= i && i < n) "AtCoder.Extra.Semigroup.Permutation.new: index boundary error" in ()) () xs + +-- | \(O(1)\) Creates a `Permutation`, without performing boundary check on input vector. +-- +-- @since 1.1.0.0 +{-# INLINE unsafeNew #-} +unsafeNew :: (HasCallStack) => VU.Vector Int -> Permutation +unsafeNew = Permutation + +-- | \(O(1)\) Creates an identity `Permutation` of length \(n\). +-- +-- @since 1.1.0.0 +{-# INLINE ident #-} +ident :: Int -> Permutation +ident = Permutation . (`VU.generate` id) + +-- | \(O(1)\) Creates a zero `Permutation` of length \(n\). It's similar to `ident`, but filled +-- with \(-1\) and invalidates corresponding slots on composition. +-- +-- @since 1.1.0.0 +{-# INLINE zero #-} +zero :: Int -> Permutation +zero n = Permutation $ VU.replicate n (-1) + +-- | \(O(1)\) Maps an index. +-- +-- @since 1.1.0.0 +{-# INLINE act #-} +act :: (HasCallStack) => Permutation -> Int -> Int +act (Permutation vec) i = case vec VG.! i of + (-1) -> i + i' -> i' + +-- | \(O(1)\) Returns the length of the internal vector. +-- +-- @since 1.1.0.0 +{-# INLINE length #-} +length :: (HasCallStack) => Permutation -> Int +length = VU.length . unPermutation + +-- | @since 1.1.0.0 +instance Semigroup Permutation where + {-# INLINE (<>) #-} + Permutation r2 <> Permutation r1 = Permutation $ VU.map f r1 + where + !_ = ACIA.runtimeAssert (VU.length r2 == VU.length r1) "AtCoder.Extra.Semigroup.Permutation.(<>): legth mismatch" + f (-1) = -1 + f i = VG.unsafeIndex r2 i diff --git a/src/AtCoder/Extra/Tree.hs b/src/AtCoder/Extra/Tree.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Tree.hs @@ -0,0 +1,207 @@ +-- | Generic tree functions. +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Tree + ( -- * Tree folding + + -- | These function are built around the three type parameters: \(w\), \(f\) and \(a\). + -- + -- - \(w\): Edge weight. + -- - \(f\): Monoid action to a vertex value. These actions are created from vertex value \(a\) + -- and edge information @(Int, w)@. + -- - \(a\): Monoid values stored at vertices. + fold, + scan, + foldReroot, + ) +where + +import Data.Functor.Identity (runIdentity) +import Data.Vector.Generic qualified as VG +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import GHC.Stack (HasCallStack) + +{-# INLINE foldImpl #-} +foldImpl :: + forall m w f a. + (HasCallStack, Monad m, VU.Unbox w) => + (Int -> VU.Vector (Int, w)) -> + (Int -> a) -> + (a -> (Int, w) -> f) -> + (f -> a -> a) -> + Int -> + (Int -> a -> m ()) -> + m a +foldImpl tree valAt toF act root memo = inner (-1) root + where + inner :: Int -> Int -> m a + inner !parent !v1 = do + let !acc0 = valAt v1 + let !v2s = VU.filter ((/= parent) . fst) $ tree v1 + !res <- VU.foldM' (\acc (!v2, !w) -> (`act` acc) . (`toF` (v1, w)) <$> inner v1 v2) acc0 v2s + memo v1 res + pure res + +-- | \(O(n)\) Folds a tree from a root vertex, also known as tree DP. +-- +-- ==== __Example__ +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import AtCoder.Extra.Tree qualified as Tree +-- >>> import Data.Semigroup (Sum (..)) +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let gr = Gr.build @(Sum Int) 5 . Gr.swapDupe $ VU.fromList [(2, 1, Sum 1), (1, 0, Sum 1), (2, 3, Sum 1), (3, 4, Sum 1)] +-- >>> type W = Sum Int -- edge weight +-- >>> type F = Sum Int -- action type +-- >>> type X = Sum Int -- vertex value +-- >>> :{ +-- let res = Tree.fold (gr `Gr.adjW`) valAt toF act 2 +-- where +-- valAt :: Int -> X +-- valAt = const $ mempty @(Sum Int) +-- toF :: X -> (Int, W) -> F +-- toF x (!_i, !dx) = x + dx +-- act :: F -> X -> X +-- act dx x = dx + x +-- in getSum res +-- :} +-- 4 +-- +-- @since 1.1.0.0 +{-# INLINE fold #-} +fold :: + (HasCallStack, VU.Unbox w) => + -- | Graph as a function. + (Int -> VU.Vector (Int, w)) -> + -- | @valAt@: Assignment of initial vertex values. + (Int -> a) -> + -- | @toF@: Converts a vertex value into an action onto a neighbor vertex. + (a -> (Int, w) -> f) -> + -- | @act@: Performs an action onto a vertex value. + (f -> a -> a) -> + -- | Root vertex. + Int -> + -- | Tree folding result from the root vertex. + a +fold tree valAt toF act root = runIdentity $ do + foldImpl tree valAt toF act root (\_ _ -> pure ()) + +-- | \(O(n)\) Folds a tree from a root vertex, also known as tree DP. The calculation process on +-- every vertex is recoreded and returned as a vector. +-- +-- ==== __Example__ +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import AtCoder.Extra.Tree qualified as Tree +-- >>> import Data.Semigroup (Sum (..)) +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let n = 5 +-- >>> let gr = Gr.build @(Sum Int) n . Gr.swapDupe $ VU.fromList [(2, 1, Sum 1), (1, 0, Sum 1), (2, 3, Sum 1), (3, 4, Sum 1)] +-- >>> type W = Sum Int -- edge weight +-- >>> type F = Sum Int -- action type +-- >>> type X = Sum Int -- vertex value +-- >>> :{ +-- let res = Tree.scan n (gr `Gr.adjW`) valAt toF act 2 +-- where +-- valAt :: Int -> X +-- valAt = const $ mempty @(Sum Int) +-- toF :: X -> (Int, W) -> F +-- toF x (!_i, !dx) = x + dx +-- act :: F -> X -> X +-- act dx x = dx + x +-- in VU.map getSum res +-- :} +-- [0,1,4,1,0] +-- +-- @since 1.1.0.0 +{-# INLINE scan #-} +scan :: + (VU.Unbox w, VG.Vector v a) => + -- | The number of vertices. + Int -> + -- | Graph as a function. + (Int -> VU.Vector (Int, w)) -> + -- | @valAt@: Assignment of initial vertex values. + (Int -> a) -> + -- | @toF@: Converts a vertex value into an action onto a neighbor vertex. + (a -> (Int, w) -> f) -> + -- | @act@: Performs an action onto a vertex value. + (f -> a -> a) -> + -- | Root vertex. + Int -> + -- | Tree scanning result from a root vertex. + v a +scan n tree acc0At toF act root = VG.create $ do + dp <- VGM.unsafeNew n + !_ <- foldImpl tree acc0At toF act root $ \v a -> do + VGM.unsafeWrite dp v a + pure dp + +-- | \(O(n)\) Folds a tree from every vertex, using the rerooting technique. +-- +-- ==== __Example__ +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import AtCoder.Extra.Tree qualified as Tree +-- >>> import Data.Semigroup (Sum (..)) +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let n = 5 +-- >>> let gr = Gr.build @(Sum Int) n . Gr.swapDupe $ VU.fromList [(2, 1, Sum 1), (1, 0, Sum 1), (2, 3, Sum 1), (3, 4, Sum 1)] +-- >>> type W = Sum Int -- edge weight +-- >>> type F = Sum Int -- action type +-- >>> type X = Sum Int -- vertex value +-- >>> :{ +-- let res = Tree.foldReroot n (gr `Gr.adjW`) valAt toF act +-- where +-- valAt :: Int -> X +-- valAt = const $ mempty @(Sum Int) +-- toF :: X -> (Int, W) -> F +-- toF x (!_i, !dx) = x + dx +-- act :: F -> X -> X +-- act dx x = dx + x +-- in VU.map getSum res +-- :} +-- [4,4,4,4,4] +-- +-- @since 1.1.0.0 +{-# INLINE foldReroot #-} +foldReroot :: + forall w f a. + (HasCallStack, VU.Unbox w, VU.Unbox a, VU.Unbox f, Monoid f) => + -- | The number of vertices. + Int -> + -- | Graph as a function. + (Int -> VU.Vector (Int, w)) -> + -- | @valAt@:Assignment of initial vertex values. + (Int -> a) -> + -- | @toF@: Converts a vertex value into an action onto a neighbor vertex. + (a -> (Int, w) -> f) -> + -- | @act@: Performs an action onto a vertex value. + (f -> a -> a) -> + -- | Tree folding result from every vertex as a root. + VU.Vector a +foldReroot n tree valAt toF act = VU.create $ do + -- Calculate tree DP for every vertex as a root: + !dp <- VUM.unsafeNew n + let reroot parent parentF v1 = do + -- TODO: when the operator is not commutative? + let !children = VU.filter ((/= parent) . fst) $ tree v1 + let !fL = VU.scanl' (\ !f (!v2, !w) -> (f <>) . (`toF` (v1, w)) $ treeDp VG.! v2) f0 children + let !fR = VU.scanr' (\(!v2, !w) !f -> (<> f) . (`toF` (v1, w)) $ treeDp VG.! v2) f0 children + + -- save + let !x1 = (parentF <> VU.last fL) `act` valAt v1 + VGM.unsafeWrite dp v1 x1 + + VU.iforM_ children $ \i2 (!v2, !w) -> do + -- composited operator excluding @v2@: + let !f1 = parentF <> (fL VG.! i2) <> (fR VG.! (i2 + 1)) + let !v1Acc = f1 `act` valAt v1 + let !f2 = toF v1Acc (v2, w) + reroot v1 f2 v2 + + reroot (-1 :: Int) f0 root0 + pure dp + where + !root0 = 0 :: Int + !f0 = mempty @f + !treeDp = scan n tree valAt toF act root0 :: VU.Vector a diff --git a/src/AtCoder/Extra/Tree/Hld.hs b/src/AtCoder/Extra/Tree/Hld.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Tree/Hld.hs @@ -0,0 +1,512 @@ +{-# LANGUAGE RecordWildCards #-} + +-- | Heavy-light decomposition is a method for partitioning a tree into segments with consecutive +-- indices. It processes various path queries in \(O(\log n)\) time. For segment tree integration +-- and monoid products, refer to the @TreeMonoid@ module. +-- +-- ==== __Overview of the internals__ +-- The following is for understanding the internals, not for using the API. Skip to the examples if +-- you want. +-- +-- ===== Original tree +-- +-- Consider a tree with arbitrary vertex order: +-- +-- @ +-- 0--8--7--3--1--2--12--13--15--14 XX: original Vertex +-- | | --: edge +-- 10--5 11--8--6 |: edge +-- | +-- 4 +-- @ +-- +-- ===== `indexHld`: Vertex -> VertexHld +-- +-- The tree vertices are reindexed with `indexHld`, where each segment is assigned consecutive +-- vertex indices: +-- +-- @ +-- 0==1==2==3==4==5==6==7==8==9 XX: VertexHld +-- | | ==: edges on the same semgent +-- 14==13 10==11==12 |: edge between different segments +-- | +-- 15 +-- @ +-- +-- Note that vertices on higher (closer to the root) segments are assigned smaller indices. This is +-- very internally very important when calculating `lca`. +-- +-- ===== `headHld`: Vertex -> Vertex +-- +-- `headHld` points the "head" vertex of each segment. It can be used for finding LCA of two +-- vertices. To find the LCA, move up to the head, go up to the parental segment's vertex and +-- repeat until the two vertices are on the same segment. +-- +-- @ +-- 0==0==0==0==0==0==0==0==0==0 XX: original Vertex +-- | | +-- 5==5 11==11==11 +-- | +-- 4 +-- @ +-- +-- `headHld` also works for identifying segments. When two vertices are on the same segment, they +-- have the same head. +-- +-- ===== `parentHld`: Vertex -> Vertex +-- +-- `parentHld` points the parental segment's vertex from a head: +-- +-- @ +-- (-1)==0==8==7==3==1==2==12==13==15 XX: original Vertex +-- | | +-- 5==8 1==11=8 +-- | +-- 5 +-- @ +-- +-- ==== __Example__ +-- Create an `Hld` for a tree: +-- +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> -- 0--1--2--3 +-- >>> -- + +-- >>> -- +--4--5 +-- >>> let n = 6 +-- >>> let tree = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (2, 3), (1, 4), (4, 5)] +-- >>> let hld = Hld.new tree +-- +-- `Hld` can process various queries in \(O(\log n)\) time: +-- +-- >>> Hld.ancestor hld 5 3 -- go up three parents from `5` +-- 0 +-- +-- >>> Hld.jump hld 5 2 3 -- go to the third vertex from `5` to `2`: +-- Just 2 +-- +-- >>> Hld.lengthBetween hld 5 3 -- get the length (the number of edges) between `5` and `3`: +-- 4 +-- +-- >>> Hld.path hld 5 3 -- get the path between `5` and `3`: +-- [5,4,1,2,3] +-- +-- Our `Hld` is rooted at @0@ vertex and subtree queries are available: +-- +-- >>> Hld.isInSubtree hld 2 3 -- `3` is in the subtree of `2` +-- True +-- +-- >>> Hld.isInSubtree hld 2 4 -- `4` is not in the subtree of `2` +-- False +-- +-- ===== Segment queries +-- Products and segment queries are primarily used by the @TreeMonoid@ module and is not intended +-- for diretct use, but here's some examples. This time the reindex by the HLD is identity: +-- +-- >>> Hld.indexHld hld +-- [0,1,2,3,4,5] +-- +-- So we can easily understand the outputs: +-- +-- >>> Hld.pathSegmentsInclusive Hld.WeightsAreOnVertices hld 5 3 +-- [(5,4),(1,3)] +-- +-- >>> Hld.pathSegmentsInclusive Hld.WeightsAreOnEdges hld 5 3 -- LCA (1) is removed +-- [(5,4),(2,3)] +-- +-- >>> Hld.subtreeSegmentInclusive hld 1 +-- (1,5) +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Tree.Hld + ( -- * Hld + Hld (..), + Vertex, + VertexHld, + + -- * Constructors + new, + newAt, + + -- * LCA + lca, + + -- * Jump + ancestor, + jump, + + -- * Path + lengthBetween, + path, + pathSegmentsInclusive, + + -- * Subtree + subtreeSegmentInclusive, + isInSubtree, + + -- * Products + WeightPolicy (..), + prod, + ) +where + +import AtCoder.Extra.Graph qualified as Gr +import AtCoder.Internal.Assert qualified as ACIA +import Control.Monad +import Control.Monad.Fix +import Control.Monad.ST +import Data.Maybe +import Data.Vector.Generic qualified as VG +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import GHC.Stack (HasCallStack) + +-- | Original graph vertex. +-- +-- @since 1.1.0.0 +type Vertex = Int + +-- | Vertex reindexed by `indexHld`. +-- +-- @since 1.1.0.0 +type VertexHld = Vertex + +-- | `Hld` partitions a tree into segments and assignes contiguous `VertexHld` for each segment. +-- +-- @since 1.1.0.0 +data Hld = Hld + { -- | The root vertex. + -- + -- @since 1.1.0.0 + rootHld :: {-# UNPACK #-} !Vertex, + -- | Maps `Vertex` to the parent `Vertex`. Returns @-1@ for the root node. + -- + -- @since 1.1.0.0 + parentHld :: !(VU.Vector Vertex), + -- | Maps `Vertex` to `VertexHld`, re-indexed vertices contiguous in each segment. + -- + -- @since 1.1.0.0 + indexHld :: !(VU.Vector VertexHld), + -- | Maps `Vertex` to the head `Vertex` of the segment. + -- + -- @since 1.1.0.0 + headHld :: !(VU.Vector Vertex), + -- | Maps `VertexHld` back to `Vertex`. Used for `ancestor` etc. + -- + -- @since 1.1.0.0 + revIndexHld :: !(VU.Vector Vertex), + -- | Maps `Vertex` to their depth from the root. Used for `jump` etc. + -- + -- @since 1.1.0.0 + depthHld :: !(VU.Vector Int), + -- | Maps `Vertex` to the subtree size. This is for subtree products. + -- + -- @since 1.1.0.0 + subtreeSizeHld :: !(VU.Vector Int) + } + deriving + ( -- | @since 1.1.0.0 + Show, + -- | @since 1.1.0.0 + Eq + ) + +-- | \(O(n)\) Creates an `Hld` with \(0\) as the root vertex. +-- +-- @since 1.1.0.0 +{-# INLINE 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 #-} +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. + -- + -- We /could/ instead record the biggest adjacent subtree vertex for each vertex, but the other + -- DFS would be harder. + let (!tree', !parent, !depths, !subtreeSize) = runST $ do + adjVec <- VU.thaw (Gr.adjCsr tree) + parent_ <- VUM.unsafeNew n + depths_ <- VUM.unsafeNew n + subtreeSize_ <- VUM.unsafeNew n + + _ <- (\f -> fix f 0 (-1) root) $ \loop depth p v1 -> do + VGM.write parent_ v1 p + VGM.write depths_ v1 depth + + (!size1, !eBig) <- + VU.foldM' + ( \(!size1, !eBig) (!e2, !v2) -> do + if v2 == p + then pure (size1, eBig) + else do + size2 <- loop (depth + 1) v1 v2 + -- NOTE: It's `>` because we should swap at least once if there's some vertex other + -- that the parent_. + pure (size1 + size2, if size1 > size2 then eBig else e2) + ) + (1 :: Int, -1) + (tree `Gr.eAdj` v1) + + -- move the biggest subtree's head to the first adjacent vertex. + -- it means the "heavy edge" or the longest segment. + when (eBig /= -1) $ do + VGM.swap adjVec eBig $ fst (VG.head (tree `Gr.eAdj` v1)) + + -- record subtree size + VGM.write subtreeSize_ v1 size1 + + pure size1 + + !vec <- VU.unsafeFreeze adjVec + (tree {Gr.adjCsr = vec},,,) + <$> VU.unsafeFreeze parent_ + <*> VU.unsafeFreeze depths_ + <*> VU.unsafeFreeze subtreeSize_ + + -- vertex -> reindexed vertex index + indices <- VUM.replicate n (-1 :: Int) + + -- vertex -> head vertex of the segment + heads <- VUM.replicate n (-1 :: Int) + + _ <- (\f -> fix f (0 :: Int) root (-1) root) $ \loop acc h p v1 -> do + -- reindex: + VGM.write indices v1 acc + let !acc' = acc + 1 + + VGM.write heads v1 h + + -- when the first vertex is within the same segment: + let (!adj1, !rest) = fromJust $ VU.uncons (tree' `Gr.adj` v1) + acc'' <- + if adj1 == p + then pure acc' + else loop acc' h v1 adj1 + + -- the others are in other segments: + VU.foldM' + ( \a v2 -> do + if v2 == p + then pure a + else loop a v2 v1 v2 + ) + acc'' + rest + + !indices' <- VU.unsafeFreeze indices + let !revIndex = VU.update (VU.replicate n (-1)) $ VU.imap (flip (,)) indices' + + Hld root parent indices' + <$> VU.unsafeFreeze heads + <*> pure revIndex + <*> pure depths + <*> pure subtreeSize + where + !n = Gr.nCsr tree + !_ = ACIA.runtimeAssert (2 * (Gr.nCsr tree - 1) == Gr.mCsr tree) "AtCoder.Extra.Hld.newAt: not a non-directed tree" + +-- | \(O(\log n)\) Calculates the lowest common ancestor of \(u\) and \(v\). +-- +-- @since 1.1.0.0 +{-# INLINE lca #-} +lca :: (HasCallStack) => Hld -> Vertex -> Vertex -> Vertex +lca Hld {..} = inner + where + inner !x !y + -- sort for easier processing + -- TODO: @case compare ix iy@ would be easier for me to understand + | ix > iy = inner y x + -- @x@ and @y@ are in other segments: + | hx /= hy = inner x $ parentHld VG.! hy + -- @x@ and @y@ are within the same segment: + -- select the smaller one, which is closer to the root and that is the LCA. + | otherwise = x + where + !ix = indexHld VG.! x + !iy = indexHld VG.! y + hx = headHld VG.! x + hy = headHld VG.! y + +-- | \(O(\log n)\) Go up \(k\) times from a vertex \(v\) to the root node. Throws an error if \(k\) +-- is bigger than the depth of \(v\). +-- +-- @since 1.1.0.0 +{-# INLINE ancestor #-} +ancestor :: (HasCallStack) => Hld -> Vertex -> Int -> Vertex +ancestor Hld {..} parent k0 = inner parent k0 + where + !_ = ACIA.runtimeAssert (0 <= k0 && k0 <= depthHld VG.! parent) $ "AtCoder.Extra.Tree.Hld.ancestor: k-th ancestor is out of the bounds (`k = " ++ show k0 ++ "`)" + inner v k + -- on this segment + | k <= iv - ihv = revIndexHld VG.! (iv - k) + -- next segment + | otherwise = inner (parentHld VG.! hv) (k - (iv - ihv + 1)) + where + iv = indexHld VG.! v + hv = headHld VG.! v + ihv = indexHld VG.! hv + +-- | \(O(\log n)\) Returns the \(k\)-th vertex of the path between \(u\) and \(v\) from \(u\). +-- Throws an error if `k` is out +-- +-- @since 1.1.0.0 +{-# INLINE jump #-} +jump :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int -> Maybe Vertex +jump hld@Hld {..} u v k + | k > lenU + lenV = Nothing + | k <= lenU = Just $ ancestor hld u k + | otherwise = Just $ ancestor hld v (lenU + lenV - k) + where + lca_ = lca hld u v + du = depthHld VG.! u + dv = depthHld VG.! v + lenU = du - depthHld VG.! lca_ + lenV = dv - depthHld VG.! lca_ + +-- | \(O(\log n)\) Returns the length of the path between \(u\) and \(v\). +-- +-- @since 1.1.0.0 +{-# INLINE lengthBetween #-} +lengthBetween :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int +lengthBetween hld@Hld {..} u v = du - dLca + dv - dLca + where + !lca_ = lca hld u v + !dLca = depthHld VG.! lca_ + !du = depthHld VG.! u + !dv = depthHld VG.! v + +-- | \(O(n)\) Returns the vertices on the path between \(u\) and \(v\). +-- +-- @since 1.1.0.0 +{-# INLINE path #-} +path :: (HasCallStack) => Hld -> Vertex -> Vertex -> [Vertex] +path hld@Hld {..} u v = concatMap expand $ pathSegmentsInclusive WeightsAreOnVertices hld u v + where + expand (!l, !r) + | l <= r = map (revIndexHld VG.!) [l .. r] + | otherwise = map (revIndexHld VG.!) [l, l - 1 .. r] + +-- | \(O(\log n)\) Decomposes a path between two vertices \(u\) and \(v\) into segments. Each +-- segment is represented as an __inclusive__ range \([u_i, v_i]\) of `VertexHLD`. +-- +-- The LCA is omitted from the returning vertices when the weight policy is set to +-- `WeightsAreOnEdges`. This is the trick to put edge weights to on vertices. +-- +-- @since 1.1.0.0 +{-# INLINE pathSegmentsInclusive #-} +pathSegmentsInclusive :: (HasCallStack) => WeightPolicy -> Hld -> Vertex -> Vertex -> [(VertexHld, VertexHld)] +pathSegmentsInclusive weightPolicy Hld {..} x0 y0 = done $ inner x0 [] y0 [] + where + isEdge = weightPolicy == WeightsAreOnEdges + done (!up, !down) = reverse up ++ down + -- @up@: bottom to top. [(max, min)] + -- @down@: top to bottom. [(min, max)] + inner :: Vertex -> [(VertexHld, VertexHld)] -> Vertex -> [(VertexHld, VertexHld)] -> ([(VertexHld, VertexHld)], [(VertexHld, VertexHld)]) + inner x up y down + | hx == hy && isEdge = case compare ix iy of + -- skip LCA on edge vertices + LT -> (up, (ix {- edge -} + 1, iy) : down) + GT -> ((ix, iy {- edge -} + 1) : up, down) + EQ -> (up, down) + | hx == hy && not isEdge = case compare ix iy of + LT -> (up, (ix, iy) : down) + _ -> ((ix, iy) : up, down) + | otherwise = case compare ix iy of + LT -> inner x up phy ((ihy, iy) : down) + GT -> inner phx ((ix, ihx) : up) y down + EQ -> error "unreachable" + where + ix, iy :: VertexHld + !ix = indexHld VG.! x + !iy = indexHld VG.! y + hx, hy :: Vertex + hx = headHld VG.! x + hy = headHld VG.! y + ihx, ihy :: VertexHld + ihx = indexHld VG.! hx + ihy = indexHld VG.! hy + phx, phy :: VertexHld + phx = parentHld VG.! hx + phy = parentHld VG.! hy + +-- | \(O(1)\) Returns a half-open interval of `VertexHld` \([\mathrm{start}, \mathrm{end})\) that +-- corresponds to the subtree segments rooted at the given @subtreeRoot@. +-- +-- @since 1.1.0.0 +{-# INLINE subtreeSegmentInclusive #-} +subtreeSegmentInclusive :: (HasCallStack) => Hld -> Vertex -> (VertexHld, VertexHld) +subtreeSegmentInclusive Hld {..} subtreeRoot = (ir, ir + sr - 1) + where + ir = indexHld VG.! subtreeRoot + sr = subtreeSizeHld VG.! subtreeRoot + +-- | \(O(1)\) Returns `True` if \(u\) is in a subtree of \(r\). +-- +-- @since 1.1.0.0 +{-# INLINE isInSubtree #-} +isInSubtree :: (HasCallStack) => Hld -> Vertex -> Vertex -> Bool +isInSubtree hld@Hld {..} r_ u = l <= iu && iu <= r + where + (!l, !r) = subtreeSegmentInclusive hld r_ + !iu = indexHld VG.! u + +-- | Represents whether weights are put on vertices or edges. +-- +-- @since 1.1.0.0 +data WeightPolicy + = -- | Weights are put on vertices. + -- + -- @since 1.1.0.0 + WeightsAreOnVertices + | -- | Weights are put on edges. + -- + -- @since 1.1.0.0 + WeightsAreOnEdges + deriving + ( -- | @since 1.1.0.0 + Eq, + -- | @since 1.1.0.0 + Show + ) + +-- | \(O(\log n f)\) Returns product of the path between \(u\) and \(v\), using the user functions +-- of time complexity \(O(f)\). +-- +-- @since 1.1.0.0 +{-# INLINE prod #-} +prod :: + (HasCallStack, Monoid mono, Monad m) => + -- | The `WeightPolicy`. + WeightPolicy -> + -- | The `Hld`. + Hld -> + -- | User function for getting products in \([u, v)\), where \(u < v\) and + -- \(\mathrm{depth}(u) < \mathrm{depth}(v)\). + (VertexHld -> VertexHld -> m mono) -> + -- | User function for getting products in \([u, v)\), where \(u < v\) and + -- \(\mathrm{depth}(u) > \mathrm{depth}(v)\). + (VertexHld -> VertexHld -> m mono) -> + -- | \(u\). + Vertex -> + -- | \(v\). + Vertex -> + -- | Product of the path between \(u\) and \(v\). + m mono +prod weightPolicy hld prodF prodB u0 v0 = do + foldM + ( \ !acc (!u, !v) -> do + !x <- + if u <= v + then prodF u $ v + 1 + else prodB v $ u + 1 + pure $! acc <> x + ) + mempty + (pathSegmentsInclusive weightPolicy hld u0 v0) diff --git a/src/AtCoder/Extra/Tree/TreeMonoid.hs b/src/AtCoder/Extra/Tree/TreeMonoid.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Tree/TreeMonoid.hs @@ -0,0 +1,297 @@ +{-# LANGUAGE RecordWildCards #-} + +-- | Integration of segment trees with the heavy-light decomposition technique. Computes monoid +-- products on a path in \(O(\log^2 n)\) time or on a subtree in \(O(\log n)\) time. +-- +-- - If vertices have weights, create a `TreeMonoid` with `fromVerts`. +-- - If edges have weights, create a tree monoid with `fromEdges`. +-- +-- ==== __(Internals) Weights on edges__ +-- +-- When vertices are unweighted and only edges have weights, treat edges as new vertices or assign +-- edge weights to the deeper vertex. +-- +-- Idea 1. Convert edges into new vertices. This is inefficient. +-- +-- @ +-- o--o--o --> o-x-o-x-o +-- @ +-- +-- Idea 2. Assign edge weight to the deeper vertex. The is the internal implementation of +-- `fromEdges` and LCAs are ignored on `prod`: +-- +-- @ +-- o +-- | <--- edge 1 +-- o <- write weight 1 here +-- | <--- edge 2 +-- o <- write weight 2 here +-- @ +-- +-- ==== __Example (1): Weights are on vertices__ +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld +-- >>> import AtCoder.Extra.Tree.TreeMonoid qualified as TM +-- >>> import Data.Semigroup (Sum (..)) +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> -- 0--1--2--3 +-- >>> -- + +-- >>> -- +--4--5 +-- >>> let n = 6 +-- >>> let tree = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (2, 3), (1, 4), (4, 5)] +-- >>> let weights = VU.generate n Sum -- vertex `i` is given weight of `i` +-- >>> let hld = Hld.new tree +-- >>> tm <- TM.fromVerts hld {- `Sum` is commutative -} Commute weights +-- >>> TM.prod tm 1 3 +-- Sum {getSum = 6} +-- +-- >>> TM.prodSubtree tm 1 +-- Sum {getSum = 15} +-- +-- >>> TM.write tm 1 $ Sum 10 +-- >>> TM.prod tm 1 3 +-- Sum {getSum = 15} +-- +-- ==== __Example (2): Weights are on edges__ +-- >>> import AtCoder.Extra.Graph qualified as Gr +-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld +-- >>> import AtCoder.Extra.Tree.TreeMonoid qualified as TM +-- >>> import Data.Semigroup (Sum (..)) +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> -- 0--1--2--3 +-- >>> -- + +-- >>> -- +--4--5 +-- >>> let n = 6 +-- >>> let edges = VU.fromList [(0, 1, Sum (1 :: Int)), (1, 2, Sum 2), (2, 3, Sum 3), (1, 4, Sum 4), (4, 5, Sum 5)] +-- >>> let tree = Gr.build n $ Gr.swapDupe edges +-- >>> let hld = Hld.new tree +-- >>> tm <- TM.fromEdges hld {- `Sum` is commutative -} Commute edges +-- >>> TM.prod tm 1 3 +-- Sum {getSum = 5} +-- +-- >>> TM.prodSubtree tm 1 +-- Sum {getSum = 14} +-- +-- >>> TM.write tm 2 $ Sum 10 +-- >>> TM.prod tm 1 3 +-- Sum {getSum = 13} +-- +-- @since 1.1.0.0 +module AtCoder.Extra.Tree.TreeMonoid + ( -- * TreeMonoid + TreeMonoid, + Vertex, + VertexHld, + Commutativity (..), + + -- * Constructors + fromVerts, + fromEdges, + + -- * Segment tree methods + + -- ** Reading + prod, + prodSubtree, + read, + + -- ** Modifications + write, + exchange, + modify, + modifyM, + ) +where + +import AtCoder.Extra.Tree.Hld qualified as Hld +import AtCoder.Internal.Assert qualified as ACIA +import AtCoder.SegTree qualified as ST +import Control.Monad +import Control.Monad.Primitive (PrimMonad, PrimState) +import Data.Monoid (Dual (..)) +import Data.Vector.Generic qualified as VG +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import GHC.Stack (HasCallStack) +import Prelude hiding (read) + +-- | Original graph vertex. +-- +-- @since 1.1.0.0 +type Vertex = Int + +-- | Vertex reindexed by `indexHld`. +-- +-- @since 1.1.0.0 +type VertexHld = Vertex + +-- | A wrapper for `Hld` getting product on paths on a tree using `Hld` and segment tree(s). +-- +-- @since 1.1.0.0 +data TreeMonoid a s = TreeMonoid + { -- | Borrowed Hld. + hldTM :: !Hld.Hld, + -- | Indicates if it's targetting commutative monoids. + commuteTM :: !Commutativity, + -- | Indicates if it's targetting edge weights (If not, it's targetting vertex weights). + weightPolicyTM :: !Hld.WeightPolicy, + -- | Segment tree for getting products upwards. + segFTM :: !(ST.SegTree s a), + -- | Segment tree for getting products downwards. Only created when the monoid is + -- `NonCommute`. + segBTM :: !(ST.SegTree s (Dual a)) + } + +-- | Represents whether a monoid is commutative or noncommutative. +-- +-- @since 1.1.0.0 +data Commutativity + = -- | Commutative: \(a \cdot b = b \cdot a\). + -- + -- @since 1.1.0.0 + Commute + | -- | Noncommutative: \(a \cdot b \neq b \cdot a\). + -- + -- @since 1.1.0.0 + NonCommute + deriving + ( -- | @since 1.1.0.0 + Eq, + -- | @since 1.1.0.0 + Show + ) + +-- | \(O(n)\) +buildImpl :: + (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => + Hld.Hld -> + Commutativity -> + Hld.WeightPolicy -> + VU.Vector a -> + m (TreeMonoid a (PrimState m)) +buildImpl hldTM commuteTM weightPolicyTM weights = do + segFTM <- ST.build weights + segBTM <- + case commuteTM of + Commute -> ST.build VU.empty + NonCommute -> ST.build $ VU.map Dual weights + pure TreeMonoid {..} + +-- | \(O(n)\) Creates a `TreeMonoid` with weights on vertices. +-- +-- @since 1.1.0.0 +fromVerts :: + (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => + -- | `Hld.Hld`. + Hld.Hld -> + -- | Whether the monoid is commutative or not. + Commutativity -> + -- | The vertex weights. + VU.Vector a -> + -- | A `TreeMonoid` with weights on vertices. + m (TreeMonoid a (PrimState m)) +fromVerts hld@Hld.Hld {indexHld} commuteTM xs_ = do + let !_ = ACIA.runtimeAssert (VU.length indexHld == VU.length xs_) $ "AtCoder.Extra.Tree.TreeMonoid.fromVerts: vertex number mismatch (`" ++ show (VU.length indexHld) ++ "` and `" ++ show (VU.length xs_) ++ "`)" + let !xs = VU.create $ do + vec <- VUM.unsafeNew $ VU.length xs_ + VU.iforM_ xs_ $ \i x -> do + VGM.write vec (indexHld VG.! i) x + pure vec + buildImpl hld commuteTM Hld.WeightsAreOnVertices xs + +-- | \(O(n)\) Creates a `TreeMonoid` with weignts on edges. The edges are not required to be +-- duplicated: only one of \((u, v, w)\) or \((v, u, w)\) is needed. +-- +-- @since 1.1.0.0 +fromEdges :: + (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => + -- | `Hld.Hld`. + Hld.Hld -> + -- | Whether the monoid is commutative or not. + Commutativity -> + -- | Input edges. + VU.Vector (Vertex, Vertex, a) -> + -- | A `TreeMonoid` with weights on edges. + m (TreeMonoid a (PrimState m)) +fromEdges hld@Hld.Hld {indexHld} commuteTM edges = do + let !xs = VU.create $ do + vec <- VUM.unsafeNew $ VU.length indexHld + VU.forM_ edges $ \(!u, !v, !w) -> do + let u' = indexHld VG.! u + let v' = indexHld VG.! v + VGM.write vec (max u' v') w + pure vec + buildImpl hld commuteTM Hld.WeightsAreOnEdges xs + +-- | \(O(\log^2 n)\) Returns the product of the path between two vertices \(u\), \(v\) (invlusive). +-- +-- @since 1.1.0.0 +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 + Commute -> Hld.prod weightPolicyTM hldTM (ST.prod segFTM) (ST.prod segFTM) u v + NonCommute -> Hld.prod weightPolicyTM hldTM (ST.prod segFTM) (\l r -> getDual <$> ST.prod segBTM l r) u v + +-- | \(O(\log n)\) Returns the product of the subtree rooted at the given `Vertex`. +-- +-- @since 1.1.0.0 +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 + case weightPolicyTM of + Hld.WeightsAreOnVertices -> ST.prod segFTM l (r + 1) + Hld.WeightsAreOnEdges -> do + -- ignore the root of the subtree, which has the minimum index among the subtree vertices + if l == r + then pure mempty + else ST.prod segFTM (l + 1) (r + 1) + +-- | \(O(1)\) Reads a `TreeMonoid` value on a `Vertex`. +-- +-- @since 1.1.0.0 +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`. +-- +-- @since 1.1.0.0 +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_ + ST.write segFTM i x + when (commuteTM == NonCommute) $ do + ST.write segBTM i $ Dual x + +-- | \(O(\log n)\) Exchanges a `TreeMonoid` value on a `Vertex`. +-- +-- @since 1.1.0.0 +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_ + !res <- ST.exchange segFTM i x + when (commuteTM == NonCommute) $ do + ST.write segBTM i $ Dual x + pure res + +-- | \(O(\log n)\) Modifies a `TreeMonoid` value on a `Vertex`. +-- +-- @since 1.1.0.0 +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_ + ST.modify segFTM f i + when (commuteTM == NonCommute) $ do + ST.modify segBTM (Dual . f . getDual) i + +-- | \(O(\log n)\) Modifies a `TreeMonoid` value on a `Vertex`. +-- +-- @since 1.1.0.0 +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_ + ST.modifyM segFTM f i + when (commuteTM == NonCommute) $ do + ST.modifyM segBTM ((Dual <$>) . f . getDual) i diff --git a/src/AtCoder/Extra/WaveletMatrix.hs b/src/AtCoder/Extra/WaveletMatrix.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/WaveletMatrix.hs @@ -0,0 +1,407 @@ +{-# LANGUAGE RecordWildCards #-} + +-- | A static Wavelet Matrix. +-- +-- ==== Notation +-- Let \(S\) be the set of values in your wavelet matrix. We use \(|S|\) to denote the number of +-- distinct values contained within this set \((|S| \lt n)\). +-- +-- @since 1.1.0.0 +module AtCoder.Extra.WaveletMatrix + ( -- * Wavelet Matrix + WaveletMatrix (..), + + -- * Constructors + build, + + -- * Access (indexing) + access, + + -- * Rank (count) + rank, + rankBetween, + + -- * Selection + + -- | ==== __Example__ + -- >>> import AtCoder.Extra.WaveletMatrix qualified as WM + -- >>> import Data.Vector.Unboxed qualified as VU + -- >>> let wm = WM.build $ VU.fromList [1,1,2,1,3] + -- >>> WM.select wm 1 + -- Just 0 + -- >>> WM.selectKth wm 2 1 + -- Just 3 + -- >>> WM.selectIn wm {- [l, r) -} 1 4 {- x -} 1 + -- Just 1 + -- >>> WM.selectKthIn wm {- [l, r) -} 1 4 {- k -} 1 {- x -} 1 + -- Just 3 + select, + selectKth, + selectIn, + selectKthIn, + + -- * Quantile (value-ordered access) + kthLargestIn, + ikthLargestIn, + kthSmallestIn, + ikthSmallestIn, + -- unsafeKthLargestIn, + -- unsafeIKthLargestIn, + -- unsafeKthSmallestIn, + -- unsafeIKthSmallestIn, + + -- * Lookup + lookupLE, + lookupLT, + lookupGE, + lookupGT, + + -- * Conversions + assocsIn, + descAssocsIn, + ) +where + +import AtCoder.Extra.Bisect +import AtCoder.Extra.WaveletMatrix.Raw qualified as Rwm +import Control.Monad +import Data.Maybe (fromJust, fromMaybe) +import Data.Vector.Algorithms.Intro qualified as VAI +import Data.Vector.Generic qualified as VG +import Data.Vector.Unboxed qualified as VU + +-- | A static Wavelet Matrix. +-- +-- @since 1.1.0.0 +data WaveletMatrix = WaveletMatrix + { -- | The internal wavelet matrix, where index compression is not automatically performed. + -- + -- @since 1.1.0.0 + rawWM :: !Rwm.RawWaveletMatrix, + -- | Index compression dictionary. + -- + -- @since 1.1.0.0 + xDictWM :: !(VU.Vector Int) + } + +-- | \(O(n \log n)\) Creates a `WaveletMatrix` from an array \(a\). +-- +-- @since 1.1.0.0 +{-# INLINE build #-} +build :: VU.Vector Int -> WaveletMatrix +build ys = + let !xDictWM = VU.uniq $ VU.modify (VAI.sortBy compare) ys + !ys' = VU.map (fromJust . lowerBound xDictWM) ys + !rawWM = Rwm.build (VG.length ys) ys' + in WaveletMatrix {..} + +-- | \(O(\log |S|)\) Returns \(a[k]\) or `Nothing` if the index is out of the bounds. Try to use the +-- original array if you can. +-- +-- @since 1.1.0.0 +{-# INLINE 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 #-} +rank :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y\) + Int -> + -- | The number of \(y\) in \([l, r)\). + Int +rank wm l r y = rankBetween wm l r y (y + 1) + +-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\). +-- +-- @since 1.1.0.0 +{-# INLINE rankBetween #-} +rankBetween :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y_1\) + Int -> + -- | \(y_2\) + Int -> + -- | The number of \(y\) in \([l, r) \times [y_1, y_2)\). + Int +rankBetween WaveletMatrix {..} l r y1 y2 + | not $ 0 <= l && l < r && r <= n = 0 + | y1' >= y2' = 0 + | otherwise = Rwm.rankBetween rawWM l r y1' y2' + where + -- Handles the case @yl@ or @yr@ is not in the dict + n = Rwm.lengthRwm rawWM + y1' = fromMaybe n (bisectR 0 (VG.length xDictWM) ((< y1) . VG.unsafeIndex xDictWM)) + y2' = maybe (-1) (+ 1) (bisectL 0 (VG.length xDictWM) ((< y2) . VG.unsafeIndex xDictWM)) + +-- | \(O(\log |S|)\) Returns the index of the first \(y\) in \(a\), or `Nothing` if \(y\) is +-- not found. +-- +-- @since 1.1.0.0 +{-# INLINE select #-} +select :: WaveletMatrix -> Int -> Maybe Int +select wm = selectKth wm 0 + +-- | \(O(\log |S|)\) Returns the index of the \(k\)-th occurrence (0-based) of \(y\), or `Nothing` +-- if no such occurrence exists. +-- +-- @since 1.1.0.0 +{-# INLINE selectKth #-} +selectKth :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(k\) + Int -> + -- | \(y\) + Int -> + -- | The index of \(k\)-th \(y\) + Maybe Int +selectKth WaveletMatrix {..} k y = do + i <- lowerBound xDictWM y + -- TODO: we don't need such an explicit branch? + let !y' = xDictWM VG.! i + guard $ y' == y + Rwm.selectKth rawWM k i + +-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the first occurrence +-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. +-- +-- @since 1.1.0.0 +{-# INLINE selectIn #-} +selectIn :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y\) + Int -> + -- | The index of the first \(y\) in \([l, r)\). + Maybe Int +selectIn wm l r = selectKthIn wm l r 0 + +-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the \(k\)-th occurrence +-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. +-- +-- @since 1.1.0.0 +{-# INLINE selectKthIn #-} +selectKthIn :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \(y\) + Int -> + -- | The index of the \(k\)-th \(y\) in \([l, r)\). + Maybe Int +selectKthIn WaveletMatrix {..} l r k y = do + i <- lowerBound xDictWM y + -- TODO: we don't need such an explicit branch? + let !y' = xDictWM VG.! i + guard $ y' == y + Rwm.selectKthIn rawWM l r k i + +-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns the index of the \(k\)-th (0-based) +-- largest value. Note that duplicated values are treated as distinct occurrences. +-- +-- @since 1.1.0.0 +{-# INLINE kthLargestIn #-} +kthLargestIn :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \(k\)-th largest \(y\) in \([l, r)\) + Maybe Int +kthLargestIn WaveletMatrix {..} l r k + | Just !y <- Rwm.kthLargestIn rawWM l r k = Just $ xDictWM VG.! y + | otherwise = Nothing + +-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns both the index and the value of the +-- \(k\)-th (0-based) largest value. Note that duplicated values are treated as distinct occurrences. +-- +-- @since 1.1.0.0 +{-# INLINE ikthLargestIn #-} +ikthLargestIn :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\) + Maybe (Int, Int) +ikthLargestIn WaveletMatrix {..} l r k + | Just (!i, !y) <- Rwm.ikthLargestIn rawWM l r k = Just (i, xDictWM VG.! y) + | otherwise = Nothing + +-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns the index of the \(k\)-th (0-based) +-- smallest value. Note that duplicated values are treated as distinct occurrences. +-- +-- @since 1.1.0.0 +{-# INLINE kthSmallestIn #-} +kthSmallestIn :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \(k\)-th largest \(y\) in \([l, r)\) + Maybe Int +kthSmallestIn WaveletMatrix {..} l r k + | Just !y <- Rwm.kthSmallestIn rawWM l r k = Just $ xDictWM VG.! y + | otherwise = Nothing + +-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns both the index and the value of the +-- \(k\)-th (0-based) smallest value. Note that duplicated values are treated as distinct occurrences. +-- +-- @since 1.1.0.0 +{-# INLINE ikthSmallestIn #-} +ikthSmallestIn :: + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\) + Maybe (Int, Int) +ikthSmallestIn WaveletMatrix {..} l r k + | Just (!i, !y) <- Rwm.ikthSmallestIn rawWM l r k = Just (i, xDictWM VG.! y) + | otherwise = Nothing + +-- | \(O(\log |S|)\) +-- +-- @since 1.1.0.0 +{-# INLINE unsafeKthSmallestIn #-} +unsafeKthSmallestIn :: WaveletMatrix -> Int -> Int -> Int -> Int +unsafeKthSmallestIn WaveletMatrix {..} l r k = + xDictWM VG.! Rwm.unsafeKthSmallestIn rawWM l r k + +-- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0]\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupLE #-} +lookupLE :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y_0\) + Int -> + -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0]\) + Maybe Int +lookupLE wm l r y0 + | r' == l' = Nothing + | rank_ == 0 = Nothing + | otherwise = Just $ unsafeKthSmallestIn wm l' r' (rank_ - 1) + where + -- clamp + l' = max 0 l + r' = min (Rwm.lengthRwm (rawWM wm)) r + rank_ = rankBetween wm l' r' minBound (y0 + 1) + +-- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0)\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupLT #-} +lookupLT :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y_0\) + Int -> + -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0)\) + Maybe Int +lookupLT wm l r y0 = lookupLE wm l r (y0 - 1) + +-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times [y_0, \infty)\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupGE #-} +lookupGE :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y_0\) + Int -> + -- | Minimum \(y\) in \([l, r) \times [y_0, \infty)\). + Maybe Int +lookupGE wm l r y0 + | r' == l' = Nothing + | rank_ >= r - l = Nothing + | otherwise = Just $ unsafeKthSmallestIn wm l r rank_ + where + -- clamp + l' = max 0 l + r' = min (Rwm.lengthRwm (rawWM wm)) r + rank_ = rankBetween wm l' r' minBound y0 + +-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupGT #-} +lookupGT :: + -- | A wavelet matrix + WaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y_0\) + Int -> + -- | Minimum \(y\) in \([l, r) \times (y_0, \infty)\) + Maybe Int +lookupGT wm l r y0 = lookupGE wm l r (y0 + 1) + +-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in range \([l, r)\) in +-- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small. +-- +-- @since 1.1.0.0 +{-# INLINE assocsIn #-} +assocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)] +assocsIn WaveletMatrix {..} l r = Rwm.assocsWith rawWM l r (xDictWM VG.!) + +-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in range \([l, r)\) in +-- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small. +-- +-- @since 1.1.0.0 +{-# INLINE 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/BitVector.hs b/src/AtCoder/Extra/WaveletMatrix/BitVector.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/WaveletMatrix/BitVector.hs @@ -0,0 +1,179 @@ +{-# LANGUAGE RecordWildCards #-} + +-- | A compact bit vector, a collection of bits that can process @rank@ in \(O(1)\) and @select@ in +-- \(O(\log n)\). +-- +-- @since 1.1.0.0 +module AtCoder.Extra.WaveletMatrix.BitVector + ( -- * Bit vector + BitVector (..), + + -- * Constructor + build, + + -- * (Internal) Word-based cumultaive sum + wordSize, + csumInPlace, + + -- * Rank + rank0, + rank1, + + -- * Select + select0, + select1, + selectKthIn0, + selectKthIn1, + ) +where + +import AtCoder.Extra.Bisect (bisectL) +import Control.Monad.Primitive (PrimMonad (PrimState)) +import Data.Bit (Bit (..)) +import Data.Bits (popCount) +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 + +-- | A compact bit vector. +-- +-- @since 1.1.0.0 +data BitVector = BitVector + { -- | Packed bits. + -- + -- @since 1.1.0.0 + bitsBv :: !(VU.Vector Bit), + -- | Cumulative sum of bits by 64 words. + -- + -- @since 1.1.0.0 + csumBv :: !(VU.Vector Int) + -- we could use Word32 for csumBv, as 2^32 is large enough + } + deriving (Eq, Show) + +-- | \(O(n)\) Creates a `BitVector`. +-- +-- @since 1.1.0.0 +{-# INLINE build #-} +build :: VU.Vector Bit -> BitVector +build bitsBv = + let csumBv = VU.create $ do + vec <- VUM.replicate ((VU.length bitsBv + 63) `div` 64 + 1) 0 + _ <- csumInPlace vec bitsBv + pure vec + in BitVector {..} + +-- | The block size \(64\) for the internal cumultaive sum in the bit vector. +-- +-- @since 1.1.0.0 +{-# INLINE wordSize #-} +wordSize :: Int +wordSize = 64 + +-- | \(O(n)\) Calculates the cumulative sum in-place for the bit vector and returns the sum. +-- +-- @since 1.1.0.0 +{-# INLINE csumInPlace #-} +csumInPlace :: + (PrimMonad m) => + -- | Cumulative sum of length \(\lceil |\mathrm{bits}| / 64 \rceil\). + VUM.MVector (PrimState m) Int -> + -- | Bits + VU.Vector Bit -> + -- | Sum of the bits + m Int +csumInPlace csum bits = do + VGM.unsafeWrite csum 0 (0 :: Int) + + -- Calcuate popCount by word. TODO: use `castToWords` for most elements + VU.ifoldM' + ( \ !acc i wordSum -> do + let !acc' = acc + wordSum + VGM.unsafeWrite csum (i + 1) acc' + pure acc' + ) + (0 :: Int) + $ VU.unfoldrExactN + (VGM.length csum - 1) + (\bits' -> (popCount (VU.take wordSize bits'), VU.drop wordSize bits')) + bits + +-- | \(O(1)\) Counts the number of \(0\) bits in the interval \([0, i)\). +-- +-- @since 1.1.0.0 +{-# INLINE rank0 #-} +rank0 :: BitVector -> Int -> Int +rank0 bv i = i - rank1 bv i + +-- | \(O(1)\) Counts the number of \(1\) bits in an interval \([0, i)\). +-- +-- @since 1.1.0.0 +{-# INLINE rank1 #-} +rank1 :: BitVector -> Int -> Int +rank1 BitVector {..} i = fromCSum + fromRest + where + -- TODO: check bounds for i? + (!nWords, !nRest) = i `divMod` wordSize + fromCSum = VG.unsafeIndex csumBv nWords + fromRest = popCount . VU.take nRest . VU.drop (nWords * wordSize) $ bitsBv + +-- | \(O(\log n)\) Returns the index of \(k\)-th \(0\) (0-based), or `Nothing` if no such bit exists. +-- +-- @since 1.1.0.0 +{-# INLINE select0 #-} +select0 :: BitVector -> Int -> Maybe Int +select0 bv = selectKthIn0 bv 0 (VG.length (bitsBv bv)) + +-- | \(O(\log n)\) Returns the index of \(k\)-th \(1\) (0-based), or `Nothing` if no such bit exists. +-- +-- @since 1.1.0.0 +{-# INLINE select1 #-} +select1 :: BitVector -> Int -> Maybe Int +select1 bv = selectKthIn1 bv 0 (VG.length (bitsBv bv)) + +-- | \(O(\log n)\) Returns the index of \(k\)-th \(0\) (0-based) in \([l, r)\), or `Nothing` if no +-- such bit exists. +-- +-- @since 1.1.0.0 +{-# INLINE selectKthIn0 #-} +selectKthIn0 :: + -- | A bit vector + BitVector -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | The index of \(k\)-th \(0\) in \([l, r)\) + Maybe Int +selectKthIn0 bv l r k + | k < 0 || nZeros <= k = Nothing + | otherwise = bisectL l r $ \i -> rank0 bv i - rankL0 < k + 1 + where + nZeros = rank0 bv r - rankL0 + rankL0 = rank0 bv l + +-- | \(O(\log n)\) Returns the index of \(k\)-th \(1\) (0-based) in \([l, r)\), or `Nothing` if no +-- such bit exists. +-- +-- @since 1.1.0.0 +{-# INLINE selectKthIn1 #-} +selectKthIn1 :: + -- | A bit vector + BitVector -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | The index of \(k\)-th \(1\) in \([l, r)\) + Maybe Int +selectKthIn1 bv l r k + | k < 0 || nOnes <= k = Nothing + | otherwise = bisectL l r $ \i -> rank1 bv i - rankL1 < k + 1 + where + nOnes = rank1 bv r - rankL1 + rankL1 = rank1 bv l diff --git a/src/AtCoder/Extra/WaveletMatrix/Raw.hs b/src/AtCoder/Extra/WaveletMatrix/Raw.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/WaveletMatrix/Raw.hs @@ -0,0 +1,651 @@ +{-# LANGUAGE RecordWildCards #-} + +-- original implementation: +-- + +-- NOTE: We could integrate cumulative sum / fenwick tree / segment tree. +-- NOTE: `topK` and `intersects` are not implemented as they are slow. + +-- | A static Wavelet Matrix without automatic index comperssion. Consider using +-- @AtCoder.Extra.WaveletMatrix@ instead. +-- +-- @since 1.1.0.0 +module AtCoder.Extra.WaveletMatrix.Raw + ( -- * RawWaveletMatrix + RawWaveletMatrix (..), + + -- * Constructors + build, + + -- * Access (indexing) + access, + + -- * rank + rankLT, + rank, + rankBetween, + + -- * Select + select, + selectKth, + selectIn, + selectKthIn, + + -- * Quantile (value-ordered access) + + -- ** Safe (total) + kthSmallestIn, + ikthSmallestIn, + kthLargestIn, + ikthLargestIn, + + -- ** Unsafe + unsafeKthSmallestIn, + unsafeIKthSmallestIn, + unsafeKthLargestIn, + unsafeIKthLargestIn, + + -- * Lookup + lookupLE, + lookupLT, + lookupGE, + lookupGT, + + -- * Conversions + assocsIn, + assocsWith, + descAssocsIn, + descAssocsInWith, + ) +where + +import AtCoder.Extra.WaveletMatrix.BitVector qualified as BV +import AtCoder.Internal.Assert qualified as ACIA +import AtCoder.Internal.Bit qualified as ACIB +import Control.Monad.ST (runST) +import Data.Bit (Bit (..)) +import Data.Bits (bit, countTrailingZeros, setBit, testBit, (.|.)) +import Data.Maybe +import Data.Vector qualified as V +import Data.Vector.Algorithms.Radix qualified as VAR +import Data.Vector.Generic qualified as VG +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import GHC.Stack (HasCallStack) + +-- | A static Wavelet Matrix without automatic index comperssion. +-- +-- @since 1.1.0.0 +data RawWaveletMatrix = RawWaveletMatrix + { -- | \(\lceil \log_2 N \rceil\). + -- + -- @since 1.1.0.0 + heightRwm :: {-# UNPACK #-} !Int, + -- | The length of the original array. + -- + -- @since 1.1.0.0 + lengthRwm :: {-# UNPACK #-} !Int, + -- | The bit matrix. Each row represents (heightRwm - 1 - iRow) bit's on/off. + -- + -- @since 1.1.0.0 + bitsRwm :: !(V.Vector BV.BitVector), + -- | The number of zeros bits in each row in the bit matrix. + -- + -- @since 1.1.0.0 + nZerosRwm :: !(VU.Vector Int) + } + deriving (Eq, Show) + +-- | \(O(n \log n)\) Creates a `RawWaveletMatrix` from a vector \(a\). +-- +-- @since 1.1.0.0 +{-# INLINE build #-} +build :: + (HasCallStack) => + -- | The number of different values in the compressed vector. + Int -> + -- | A compressed vector + VU.Vector Int -> + -- | A wavelet matrix + RawWaveletMatrix +build nx xs + | nx < 0 = error "AtCoder.Extra.WaveletMatrix.Raw.build: given negative `n`" + | otherwise = runST $ do + -- TODO: less mutable variables + orgBits <- VUM.replicate (lengthRwm * heightRwm) $ Bit False + orgCsum <- VUM.replicate (lenCSum * heightRwm) (0 :: Int) + nZeros <- VUM.unsafeNew heightRwm + + -- views by row over the contiguous memory: + let !bits = V.unfoldrExactN heightRwm (VUM.splitAt lengthRwm) orgBits + let !csums = V.unfoldrExactN heightRwm (VUM.splitAt lenCSum) orgCsum + + -- the vector will be sorted by bits. + vec <- VU.thaw xs + V.izipWithM_ + ( \iRow bitVec csum -> do + let !iBit = heightRwm - 1 - iRow + vec' <- VU.unsafeFreeze vec + VU.iforM_ vec' $ \i x -> do + VGM.unsafeWrite bitVec i . Bit $ testBit x iBit + + -- csum. + VGM.unsafeWrite csum 0 (0 :: Int) + bitVec' <- VU.unsafeFreeze bitVec + + -- get popCount by word. TODO: use `castToWords` for most elements + nOnes <- BV.csumInPlace csum bitVec' + VGM.unsafeWrite nZeros iRow (lengthRwm - nOnes) + + -- preform a stable sort by the bit: + VAR.sortBy 2 2 (\_ x -> fromEnum (testBit x iBit)) vec + ) + bits + csums + nZerosRwm <- VU.unsafeFreeze nZeros + bits' <- V.unfoldrExactN heightRwm (VU.splitAt lengthRwm) <$> VU.unsafeFreeze orgBits + csums' <- V.unfoldrExactN heightRwm (VU.splitAt lenCSum) <$> VU.unsafeFreeze orgCsum + let !bitsRwm = V.zipWith BV.BitVector bits' csums' + pure $ RawWaveletMatrix {..} + where + !lengthRwm = VG.length xs + !lenCSum = (lengthRwm + BV.wordSize - 1) `div` BV.wordSize + 1 -- +1 for the zero + !heightRwm = countTrailingZeros $ ACIB.bitCeil nx + +-- | \(O(\log |S|)\) Returns \(a[k]\) or `Nothing` if the index is out of the bounds. Try to use the +-- original array if you can. +-- +-- @since 1.1.0.0 +{-# INLINE access #-} +access :: RawWaveletMatrix -> Int -> Maybe Int +access RawWaveletMatrix {..} i0 + | ACIA.testIndex i0 lengthRwm = + let (!_, !res) = + V.ifoldl' + ( \(!i, !acc) !iRow !bits -> + let Bit !goRight = VG.unsafeIndex (BV.bitsBv bits) i + !i' + | goRight = BV.rank1 bits i + VG.unsafeIndex nZerosRwm iRow + | otherwise = BV.rank0 bits i + !acc' + | goRight = setBit acc (heightRwm - 1 - iRow) + | otherwise = acc + in (i', acc') + ) + (i0, 0) + bitsRwm + in Just res + | otherwise = Nothing + +-- | \(O(\log |A|)\) Goes down the wavelet matrix for collecting the kth smallest value. +-- +-- @since 1.1.0.0 +{-# INLINE 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 + -- It's binary search over the value range. In each row, we'll focus on either 0 bit values or + -- 1 bit values in [l, r) and update the range to [l', r'). + step (!acc, !l, !r, !k) !iRow !bits + -- `r0 - l0`, the number of zeros in [l, r), is bigger than or equal to k: + -- Go left. + | k < r0 - l0 = (acc, l0, r0, k) + -- Go right. + | otherwise = + let !acc' = acc .|. bit (heightRwm - 1 - iRow) + !nZeros = VG.unsafeIndex nZerosRwm iRow + -- every zero bits come to the left after the move. + !l' = l + nZeros - l0 -- add the number of zeros in [0, l) + !r' = r + nZeros - r0 -- add the number of zeros in [0, r) + !k' = k - (r0 - l0) -- `r0 - l0` zeros go left + in (acc', l', r', k') + where + !l0 = BV.rank0 bits l + !r0 = BV.rank0 bits r + +-- | \(O(\log |A|)\) Goes up the wavelet matrix for collecting the value \(x\). +-- +-- @since 1.1.0.0 +{-# INLINE goUp #-} +goUp :: RawWaveletMatrix -> Int -> Int -> Maybe Int +goUp RawWaveletMatrix {..} i0 x = + V.ifoldM' + ( \ !i !iBit !bits -> + if testBit x iBit + then BV.select1 bits $ i - nZerosRwm VG.! (heightRwm - 1 - iBit) + else BV.select0 bits i + ) + i0 + (V.reverse bitsRwm) + +-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [0, y_0)\). +-- +-- @since 1.1.0.0 +{-# INLINE 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. + | xr >= bit heightRwm = r'_ - l'_ + | xr <= 0 = 0 + | r'_ <= l'_ = 0 + | otherwise = + let (!res, !_, !_) = V.ifoldl' step (0, l'_, r'_) bitsRwm + in res + where + -- clamp + l'_ = max 0 l_ + r'_ = min lengthRwm r_ + -- [l, r) + step (!acc, !l, !r) !iRow !bits = + let !b = testBit xr (heightRwm - 1 - iRow) + !l0 = BV.rank0 bits l + !r0 = BV.rank0 bits r + in if b + then (acc + r0 - l0, l - l0 + VG.unsafeIndex nZerosRwm iRow, r - r0 + VG.unsafeIndex nZerosRwm iRow) + else (acc, l0, r0) + +-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\). +-- +-- @since 1.1.0.0 +{-# INLINE rank #-} +rank :: + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y\) + Int -> + -- | The number of \(y\) in \([l, r)\). + Int +rank wm l r x = rankBetween wm l r x (x + 1) + +-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\). +-- +-- @since 1.1.0.0 +{-# INLINE rankBetween #-} +rankBetween :: + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y_1\) + Int -> + -- | \(y_2\) + Int -> + -- | The number of \(y\) in \([l, r) \times [y_1, y_2)\). + Int +rankBetween wm l r lx rx = rankLT wm l r rx - rankLT wm l r lx + +-- | \(O(\log |S|)\) Returns the index of the first \(y\) in the sequence, or `Nothing` if \(y\) is +-- not found. +-- +-- @since 1.1.0.0 +{-# INLINE select #-} +select :: RawWaveletMatrix -> Int -> Maybe Int +select wm = selectKth wm 0 + +-- | \(O(\log |S|)\) Returns the index of the \(k\)-th occurrence (0-based) of \(y\), or `Nothing` +-- if no such occurrence exists. +-- +-- @since 1.1.0.0 +{-# INLINE selectKth #-} +selectKth :: + RawWaveletMatrix -> + -- | \(k\) + Int -> + -- | \(y\) + Int -> + -- | The index of \(k\)-th \(y\) + Maybe Int +selectKth wm = selectKthIn wm 0 (lengthRwm wm) + +-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the first occurrence +-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. +-- +-- @since 1.1.0.0 +{-# INLINE selectIn #-} +selectIn :: + -- | A wavelet matrix + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | The index of the first \(y\) in \([l, r)\). + Maybe Int +selectIn wm = selectKthIn wm 0 + +-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the \(k\)-th occurrence +-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. +-- +-- @since 1.1.0.0 +{-# INLINE selectKthIn #-} +selectKthIn :: + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \(y\) + Int -> + -- | The index of the \(k\)-th \(y\) in \([l, r)\). + Maybe Int +selectKthIn wm@RawWaveletMatrix {..} l_ r_ k x + | not (0 <= x && x < lengthRwm && 0 <= k && k < lengthRwm) = Nothing + | l'_ < r'_ = inner + | otherwise = Nothing + where + -- clamp + l'_ = max 0 l_ + r'_ = min lengthRwm r_ + inner :: Maybe Int + inner + | rEnd <= lEnd + k = Nothing + -- go up + | otherwise = goUp wm (lEnd + k) x + where + -- TODO: replace with goDown + -- Go down. Gets the [l, r) range of @x@ in the last array. + (!lEnd, !rEnd) = + V.ifoldl' + ( \(!l, !r) !iRow !bits -> + let !l0 = BV.rank0 bits l + !r0 = BV.rank0 bits r + in if testBit x (heightRwm - 1 - iRow) + then (l + nZerosRwm VG.! iRow - l0, r + nZerosRwm VG.! iRow - r0) + else (l0, r0) + ) + (l'_, r'_) + bitsRwm + +-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the \(k\)-th (0-based) +-- largest value. Note that duplicated values are counted as distinct occurrences. +-- +-- @since 1.1.0.0 +{-# INLINE kthLargestIn #-} +kthLargestIn :: + -- | A wavelet matrix + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \(k\)-th largest \(y\) in \([l, r)\) + Maybe Int +kthLargestIn wm l r k + | k < 0 || k >= r - l = Nothing + | 0 <= l && l < r && r <= lengthRwm wm = Just $ unsafeKthLargestIn wm l r k + | otherwise = Nothing + +-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns both the index and the value of the +-- \(k\)-th (0-based) largest value. Note that duplicated values are counted as distinct occurrences. +-- +-- @since 1.1.0.0 +{-# INLINE ikthLargestIn #-} +ikthLargestIn :: + -- | A wavelet matrix + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\) + Maybe (Int, Int) +ikthLargestIn wm l r k + | k < 0 || k >= r - l = Nothing + | 0 <= l && l < r && r <= lengthRwm wm = Just $ unsafeIKthLargestIn wm l r k + | otherwise = Nothing + +-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the \(k\)-th (0-based) +-- smallest value. Note that duplicated values are counted as distinct occurrences. +-- +-- @since 1.1.0.0 +{-# INLINE kthSmallestIn #-} +kthSmallestIn :: + -- | A wavelet matrix + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \(k\)-th largest \(y\) in \([l, r)\) + Maybe Int +kthSmallestIn wm l r k + | k < 0 || k >= r - l = Nothing + | 0 <= l && l < r && r <= lengthRwm wm = Just $ unsafeKthSmallestIn wm l r k + | otherwise = Nothing + +-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns both the index and the value of the +-- \(k\)-th (0-based) smallest value. Note that duplicated values are counted as distinct occurrences. +-- +-- @since 1.1.0.0 +{-# INLINE ikthSmallestIn #-} +ikthSmallestIn :: + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(k\) + Int -> + -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\) + Maybe (Int, Int) +ikthSmallestIn wm l r k + | k < 0 || k >= r - l = Nothing + | 0 <= l && l < r && r <= lengthRwm wm = Just $ unsafeIKthSmallestIn wm l r k + | otherwise = Nothing + +-- | \(O(\log a)\) Returns \(k\)-th (0-based) biggest number in \([l, r)\). Note that duplicated +-- values are counted as distinct occurrences. +-- +-- @since 1.1.0.0 +{-# INLINE 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 #-} +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 #-} +unsafeKthSmallestIn :: RawWaveletMatrix -> Int -> Int -> Int -> Int +unsafeKthSmallestIn wm l_ r_ k_ = + let (!x, !_, !_, !_) = goDown wm l_ r_ k_ + in x + +-- | \(O(\log a)\) +-- +-- @since 1.1.0.0 +{-# INLINE unsafeIKthSmallestIn #-} +unsafeIKthSmallestIn :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int) +unsafeIKthSmallestIn wm l_ r_ k_ = + let (!x, !l, !_, !k) = goDown wm l_ r_ k_ + !i' = fromJust $ goUp wm (l + k) x + in (i', x) + +-- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0]\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupLE #-} +lookupLE :: + -- | A wavelet matrix + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y_0\) + Int -> + -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0]\) + Maybe Int +lookupLE wm l r x + | r' == l' = Nothing + | rank_ == 0 = Nothing + | otherwise = Just $ unsafeKthSmallestIn wm l' r' (rank_ - 1) + where + -- clamp + l' = max 0 l + r' = min (lengthRwm wm) r + rank_ = rankBetween wm l r minBound (x + 1) + +-- | \(O(\log a)\) Finds the maximum \(x\) in \([l, r)\) s.t. \(x_{0} \lt x\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupLT #-} +lookupLT :: + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(x\) + Int -> + -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0)\) + Maybe Int +lookupLT wm l r x = lookupLE wm l r (x - 1) + +-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times [y_0, \infty)\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupGE #-} +lookupGE :: + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y_0\) + Int -> + -- | Minimum \(y\) in \([l, r) \times [y_0, \infty)\). + Maybe Int +lookupGE wm l r x + | r' == l' = Nothing + | rank_ >= r' - l' = Nothing + | otherwise = + Just $ unsafeKthSmallestIn wm l' r' rank_ + where + -- clamp + l' = max 0 l + r' = min (lengthRwm wm) r + rank_ = rankBetween wm l' r' minBound x + +-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\). +-- +-- @since 1.1.0.0 +{-# INLINE lookupGT #-} +lookupGT :: + RawWaveletMatrix -> + -- | \(l\) + Int -> + -- | \(r\) + Int -> + -- | \(y_0\) + Int -> + -- | Minimum \(y\) in \([l, r) \times (y_0, \infty)\) + Maybe Int +lookupGT wm l r x = lookupGE wm l r (x + 1) + +-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in range \([l, r)\) in +-- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small. +-- +-- @since 1.1.0.0 +{-# INLINE 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 #-} +assocsWith :: RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)] +assocsWith RawWaveletMatrix {..} l_ r_ f + | l'_ < r'_ = inner (0 :: Int) (0 :: Int) l'_ r'_ [] + | otherwise = [] + where + -- clamp + l'_ = max 0 l_ + r'_ = min lengthRwm r_ + -- DFS. [l, r) + inner !acc iRow !l !r res + | iRow >= heightRwm = + let !n = r - l + !acc' = f acc + in (acc', n) : res + | otherwise = do + let !bits = bitsRwm VG.! iRow + !l0 = BV.rank0 bits l + !r0 = BV.rank0 bits r + !nZeros = nZerosRwm VG.! iRow + -- go right (visit bigger values first) + !l' = l + nZeros - l0 + !r' = r + nZeros - r0 + !res' + | l' < r' = inner (acc .|. bit (heightRwm - 1 - iRow)) (iRow + 1) l' r' res + | otherwise = res + !res'' + -- go left + | l0 < r0 = inner acc (iRow + 1) l0 r0 res' + | otherwise = res' + in res'' + +-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in range \([l, r)\) in +-- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small. +-- +-- @since 1.1.0.0 +{-# INLINE 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 #-} +descAssocsInWith :: RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)] +descAssocsInWith RawWaveletMatrix {..} l_ r_ f + | l'_ < r'_ = inner (0 :: Int) (0 :: Int) l'_ r'_ [] + | otherwise = [] + where + -- clamp + l'_ = max 0 l_ + r'_ = min lengthRwm r_ + -- DFS. [l, r) + inner !acc iRow !l !r res + | iRow >= heightRwm = + let !n = r - l + !acc' = f acc + in (acc', n) : res + | otherwise = do + let !bits = bitsRwm VG.! iRow + !l0 = BV.rank0 bits l + !r0 = BV.rank0 bits r + !nZeros = nZerosRwm VG.! iRow + !res' + -- go left + | l0 < r0 = inner acc (iRow + 1) l0 r0 res + | otherwise = res + -- go right (visit bigger values first) + !l' = l + nZeros - l0 + !r' = r + nZeros - r0 + !res'' + | l' < r' = inner (acc .|. bit (heightRwm - 1 - iRow)) (iRow + 1) l' r' res' + | otherwise = res' + in res'' diff --git a/src/AtCoder/Extra/WaveletMatrix2d.hs b/src/AtCoder/Extra/WaveletMatrix2d.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/WaveletMatrix2d.hs @@ -0,0 +1,277 @@ +{-# LANGUAGE RecordWildCards #-} + +-- | A 2D, static wavelet matrix with segment tree, that can handle point add and rectangle sum +-- queries. Points cannot be added after construction, but monoid values in each point can be +-- modified later. +-- +-- ==== __Example__ +-- Create a `WaveletMatrix2d` with initial vertex values: +-- +-- >>> import AtCoder.Extra.WaveletMatrix2d qualified as WM +-- >>> import Data.Semigroup (Sum (..)) +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> -- 8 9 10 11 +-- >>> -- 4 5 6 7 +-- >>> -- 0 1 2 3 +-- >>> wm <- WM.build negate $ VU.generate 12 $ \i -> let (!y, !x) = i `divMod` 4 in (x, y, Sum i) +-- +-- Read the value at \(x = 2, y = 1\): +-- +-- >>> WM.read wm (2, 1) +-- Sum {getSum = 6} +-- +-- Other segment tree methods are also available, but in 2D: +-- +-- >>> WM.allProd wm -- (0 + 11) * 12 / 2 = 66 +-- Sum {getSum = 66} +-- +-- >>> WM.prod wm {- x -} 1 3 {- y -} 0 3 -- 1 + 2 + 5 + 6 + 9 + 10 +-- Sum {getSum = 33} +-- +-- >>> WM.modify wm (+ 2) (1, 1) +-- >>> WM.prod wm {- x -} 1 3 {- y -} 0 3 -- 1 + 2 + 7 + 6 + 9 + 10 +-- Sum {getSum = 35} +-- +-- >>> WM.write wm (1, 1) $ Sum 0 +-- >>> WM.prod wm {- x -} 1 3 {- y -} 0 3 -- 1 + 2 + 0 + 6 + 9 + 10 +-- Sum {getSum = 28} +module AtCoder.Extra.WaveletMatrix2d + ( -- * Wavelet matrix 2D + WaveletMatrix2d (..), + + -- * Counstructor + new, + build, + + -- * Segment tree methods + read, + write, + modify, + prod, + prodMaybe, + allProd, + -- wavelet matrix methods could be implemented, too + ) +where + +import AtCoder.Extra.Bisect (bisectR, lowerBound) +import AtCoder.Extra.WaveletMatrix.BitVector qualified as BV +import AtCoder.Extra.WaveletMatrix.Raw qualified as Rwm +import AtCoder.Internal.Assert qualified as ACIA +import AtCoder.SegTree qualified as ST +import Control.Monad.Primitive (PrimMonad, PrimState) +import Data.Bit (Bit (..)) +import Data.Bits (Bits (testBit)) +import Data.Maybe (fromJust, fromMaybe) +import Data.Vector qualified as V +import Data.Vector.Algorithms.Intro qualified as VAI +import Data.Vector.Generic qualified as VG +import Data.Vector.Unboxed qualified as VU +import GHC.Stack (HasCallStack) +import Prelude hiding (read) + +-- NOTE: There are many possible improvements. +-- - Use cumulative sum or fenwick tree instead for the speed. +-- - The inverse operator is not actually required. +-- - Wavelet matrix methods such as `rank` can be implemented +-- - `maxRight` can be implemented. + +-- | Segment Tree on Wavelet Matrix: points on a 2D plane and rectangle products. +data WaveletMatrix2d s a = WaveletMatrix2d + { -- | The wavelet matrix that represents points on a 2D plane. + rawWmWm2d :: !Rwm.RawWaveletMatrix, + -- | (x, y) index compression dictionary. + xyDictWm2d :: !(VU.Vector (Int, Int)), + -- | y index compression dictionary. + yDictWm2d :: !(VU.Vector Int), + -- | The segment tree of the weights of the points in the order of `xyDictWm2d`. + segTreesWm2d :: !(V.Vector (ST.SegTree s a)), + -- | The inverse operator of the interested monoid. + invWm2d :: !(a -> a) + } + +-- | \(O(n \log n)\) Creates a `WaveletMatrix2d` with `mempty` as the initial monoid +-- values for each point. +{-# INLINE new #-} +new :: + (PrimMonad m, Monoid a, VU.Unbox a) => + -- | Inverse operator of the monoid + (a -> a) -> + -- | Input points + VU.Vector (Int, Int) -> + -- | A 2D wavelet matrix + m (WaveletMatrix2d (PrimState m) a) +new invWm2d xys = do + let n = VG.length xys + let xyDictWm2d = VU.uniq . VU.modify (VAI.sortBy compare) $ xys + let (!_, !ys) = VU.unzip xys + let yDictWm2d = VU.uniq $ VU.modify (VAI.sortBy compare) ys + -- REMARK: Be sure to use `n + 1` because the product function cannot handle the case + -- `yUpper` is `2^{height}`. + let (!_, !ysInput) = VU.unzip xyDictWm2d + let rawWmWm2d = Rwm.build (n + 1) $ VU.map (fromJust . lowerBound yDictWm2d) ysInput + segTreesWm2d <- V.replicateM (Rwm.heightRwm rawWmWm2d) (ST.new n) + pure WaveletMatrix2d {..} + +-- | \(O(n \log n)\) Creates a `WaveletMatrix2d` with wavelet matrix with segment tree +-- with initial monoid values. Monoids on a duplicate point are accumulated with `(<>)`. +{-# INLINE build #-} +build :: + (PrimMonad m, Monoid a, VU.Unbox a) => + -- | Inverse operator of the monoid + (a -> a) -> + -- | Input points with initial values + VU.Vector (Int, Int, a) -> + -- | A 2D wavelet matrix + m (WaveletMatrix2d (PrimState m) a) +build invWm2d xysw = do + let (!xs, !ys, !_) = VU.unzip3 xysw + wm <- new invWm2d $ VU.zip xs ys + -- not the fastest implementation though + VU.forM_ xysw $ \(!x, !y, !w) -> do + modify wm (<> w) (x, y) + pure wm + +-- | \(O(1)\) Returns the monoid value at \((x, y)\). +{-# INLINE read #-} +read :: (HasCallStack, VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> (Int, Int) -> m a +read WaveletMatrix2d {..} (!x, !y) = do + ST.read (V.head segTreesWm2d) . fromJust $ lowerBound xyDictWm2d (x, y) + +-- | \(O(\log^2 n)\) Writes the monoid value at \((x, y)\). Access to unknown points are undefined. +{-# INLINE write #-} +write :: (HasCallStack, Monoid a, VU.Unbox a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> (Int, Int) -> a -> m () +write WaveletMatrix2d {..} (!x, !y) v = do + let !i_ = fromJust $ lowerBound xyDictWm2d (x, y) + V.ifoldM'_ + ( \i iRow (!bits, !seg) -> do + let !i0 = BV.rank0 bits i + let !i' + | unBit $ VG.unsafeIndex (BV.bitsBv bits) i = + i + Rwm.nZerosRwm rawWmWm2d VG.! iRow - i0 + | otherwise = i0 + ST.write seg i' v + pure i' + ) + i_ + $ V.zip (Rwm.bitsRwm rawWmWm2d) segTreesWm2d + +-- | \(O(\log^2 n)\) Modifies the monoid value at \((x, y)\). Access to unknown points are +-- undefined. +{-# INLINE modify #-} +modify :: (HasCallStack, Monoid a, VU.Unbox a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> (a -> a) -> (Int, Int) -> m () +modify WaveletMatrix2d {..} f (!x, !y) = do + let !i_ = fromJust $ lowerBound xyDictWm2d (x, y) + V.ifoldM'_ + ( \i iRow (!bits, !seg) -> do + let !i0 = BV.rank0 bits i + let !i' + | unBit $ VG.unsafeIndex (BV.bitsBv bits) i = + i + Rwm.nZerosRwm rawWmWm2d VG.! iRow - i0 + | otherwise = i0 + ST.modify seg f i' + pure i' + ) + i_ + $ V.zip (Rwm.bitsRwm rawWmWm2d) segTreesWm2d + +-- | \(O(\log^2 n)\) Returns the monoid product in \([l, r) \times [y_1, y_2)\). +{-# INLINE prod #-} +prod :: (HasCallStack, VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a +prod wm@WaveletMatrix2d {..} !xl !xr !yl !yr + | xl' >= xr' || yl' >= yr' = pure mempty + | otherwise = unsafeProd wm xl' xr' yl' yr' + where + (!xDict, !_) = VU.unzip xyDictWm2d + -- NOTE: clamping here! + xl' = fromMaybe 0 $ bisectR 0 (VG.length xDict) $ (< xl) . VG.unsafeIndex xDict + xr' = fromMaybe (VG.length xDict) $ bisectR 0 (VG.length xDict) $ (< xr) . VG.unsafeIndex xDict + yl' = fromMaybe 0 $ bisectR 0 (VG.length yDictWm2d) $ (< yl) . VG.unsafeIndex yDictWm2d + yr' = fromMaybe (VG.length yDictWm2d) $ bisectR 0 (VG.length yDictWm2d) $ (< yr) . VG.unsafeIndex yDictWm2d + !_ = ACIA.checkInterval "AtCoder.Extra.WaveletMatrix.SegTree.prod (compressed x)" xl' xr' (VG.length xDict) + !_ = ACIA.checkInterval "AtCoder.Extra.WaveletMatrix.SegTree.prod (compressed y)" yl' yr' (VG.length yDictWm2d) + +-- | \(O(\log^2 n)\) Returns the monoid product in \([l, r) \times [y_1, y_2)\). Returns `Nothing` for invalid +-- intervals. +{-# INLINE prodMaybe #-} +prodMaybe :: (VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> Int -> Int -> Int -> Int -> m (Maybe a) +prodMaybe wm@WaveletMatrix2d {..} !xl !xr !yl !yr + | not (ACIA.testInterval xl' xr' (VG.length xDict)) = pure Nothing + | not (ACIA.testInterval yl' yr' (VG.length yDictWm2d)) = pure Nothing + | xl' >= xr' || yl' >= yr' = pure $ Just mempty + | otherwise = Just <$> unsafeProd wm xl' xr' yl' yr' + where + (!xDict, !_) = VU.unzip xyDictWm2d + -- NOTE: clamping here! + xl' = fromMaybe 0 $ bisectR 0 (VG.length xDict) $ (< xl) . VG.unsafeIndex xDict + xr' = fromMaybe (VG.length xDict) $ bisectR 0 (VG.length xDict) $ (< xr) . VG.unsafeIndex xDict + yl' = fromMaybe 0 $ bisectR 0 (VG.length yDictWm2d) $ (< yl) . VG.unsafeIndex yDictWm2d + yr' = fromMaybe (VG.length yDictWm2d) $ bisectR 0 (VG.length yDictWm2d) $ (< yr) . VG.unsafeIndex yDictWm2d + +-- | \(O(\log^2 n)\) Return the monoid product of all of the points in the wavelet matrix. +{-# INLINE allProd #-} +allProd :: (HasCallStack, VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> m a +allProd WaveletMatrix2d {..} = do + -- ST.allProd (V.last segTreesWm2d) + ST.allProd (V.head segTreesWm2d) + +-- | \(O(\log^2 n)\) The input is compressed indices. +{-# INLINE unsafeProd #-} +unsafeProd :: (VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a +unsafeProd wm xl' xr' yl' yr' = do + sR <- prodLT wm xl' xr' yr' + sL <- prodLT wm xl' xr' yl' + pure $! sR <> invWm2d wm sL + +-- | \(O(\log^2 n)\) +{-# INLINE prodLT #-} +prodLT :: (Monoid a, VU.Unbox a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> Int -> Int -> Int -> m a +prodLT WaveletMatrix2d {..} !l_ !r_ yUpper = do + (!res, !_, !_) <- do + V.ifoldM' + ( \(!acc, !l, !r) !iRow (!bits, !seg) -> do + let !l0 = BV.rank0 bits l + !r0 = BV.rank0 bits r + -- REMARK: The function cannot handle the case yUpper = N = 2^i. See the constructor for + -- how it's handled and note that l_ and r_ are compressed indices. + if testBit yUpper (Rwm.heightRwm rawWmWm2d - 1 - iRow) + then do + !acc' <- (acc <>) <$> ST.prod seg l0 r0 + let !l' = l + Rwm.nZerosRwm rawWmWm2d VG.! iRow - l0 + let !r' = r + Rwm.nZerosRwm rawWmWm2d VG.! iRow - r0 + pure (acc', l', r') + else do + pure (acc, l0, r0) + ) + (mempty, l_, r_) + $ V.zip (Rwm.bitsRwm rawWmWm2d) segTreesWm2d + pure res + +-- -- | \(O(\log n)\) Restore the original \(x\) coordinate from a compressed one. Access to unknown +-- -- points are undefined. +-- {-# INLINE indexX #-} +-- indexX :: (HasCallStack) => WaveletMatrix2d s a -> Int -> Int +-- indexX WaveletMatrix2d {xyDictWm2d} x = maybe err (VG.unsafeIndex xDict) $ lowerBound xDict x +-- where +-- (!xDict, !_) = VU.unzip xyDictWm2d +-- err = error $ "AtCoder.Extra.WaveletMatirx.SegTree.indexX: cannot index x (`" ++ show x ++ "`)" + +-- -- | \(O(\log n)\) Restore the original \(y\) coordinate from a compressed one. Access to unknown +-- -- points are undefined. +-- {-# INLINE indexY #-} +-- indexY :: (HasCallStack) => WaveletMatrix2d s a -> Int -> Int +-- indexY WaveletMatrix2d {yDictWm2d} y = maybe err (VG.unsafeIndex yDictWm2d) $ lowerBound yDictWm2d y +-- where +-- err = error $ "AtCoder.Extra.WaveletMatirx.SegTree.indexY: cannot index y (`" ++ show y ++ "`)" + +-- -- | \(O(\log n)\) Restore the original \((x, y)\) coordinates from a compressed one. Access to +-- -- unknown points are undefined. +-- {-# INLINE indexXY #-} +-- indexXY :: (HasCallStack) => WaveletMatrix2d s a -> Int -> Int -> (Int, Int) +-- indexXY WaveletMatrix2d {xyDictWm2d} x y = maybe err (VG.unsafeIndex xyDictWm2d) $ lowerBound xyDictWm2d (x, y) +-- where +-- err = error $ "AtCoder.Extra.WaveletMatirx.SegTree.indexXY: cannot index (x, y) `" ++ show (x, y) ++ "`" + +-- {-# INLINE assocsWith #-} +-- assocsWith :: WaveletMatrix -> (Int -> Int) -> [(Int, Int)] +-- assocsWith WaveletMatrix {..} l_ r_ f diff --git a/src/AtCoder/FenwickTree.hs b/src/AtCoder/FenwickTree.hs --- a/src/AtCoder/FenwickTree.hs +++ b/src/AtCoder/FenwickTree.hs @@ -1,13 +1,13 @@ {-# LANGUAGE RecordWildCards #-} --- | Fenwick tree, also known as binary index tree. Given an array of length \(n\), it processes the --- following queries in \(O(\log n)\) time. +-- | A Fenwick tree, also known as binary indexed tree. Given an array of length \(n\), it processes +-- the following queries in \(O(\log n)\) time. -- -- - Updating an element -- - Calculating the sum of the elements of an interval -- -- ==== __Example__ --- You can create a `FenwickTree` with `new`: +-- Create a `FenwickTree` with `new`: -- -- >>> import AtCoder.FenwickTree qualified as FT -- >>> ft <- FT.new @_ @Int 4 -- [0, 0, 0, 0] @@ -24,7 +24,7 @@ -- >>> FT.sum ft 0 3 -- 6 -- --- You can create a `FenwickTree` with initial values using `build`: +-- Create a `FenwickTree` with initial values using `build`: -- -- >>> ft <- FT.build @_ @Int $ VU.fromList [3, 0, 3, 0] -- >>> FT.add ft 1 2 -- [3, 2, 3, 0] @@ -40,7 +40,7 @@ new, build, - -- * Modifying the Fenwick tree + -- * Adding add, -- * Accessor @@ -60,7 +60,7 @@ import GHC.Stack (HasCallStack) import Prelude hiding (sum) --- | Fenwick tree. +-- | A Fenwick tree. -- -- @since 1.0.0.0 data FenwickTree s a = FenwickTree @@ -122,7 +122,7 @@ VGM.modify dataFt (+ x) (p - 1) loop $! p + (p .&. (-p)) --- | \(O(\log n)\) Calculates the sum in half-open range @[0, r)@. +-- | \(O(\log n)\) Calculates the sum in a half-open interval @[0, r)@. -- -- @since 1.0.0.0 {-# INLINE prefixSum #-} @@ -135,7 +135,7 @@ dx <- VGM.read dataFt (r - 1) inner (acc + dx) (r - r .&. (-r)) --- | Calculates the sum in half-open range \([l, r)\). +-- | Calculates the sum in a half-open interval \([l, r)\). -- -- ==== Constraints -- - \(0 \leq l \leq r \leq n\) @@ -150,8 +150,8 @@ | not (ACIA.testInterval l r nFt) = ACIA.errorInterval "AtCoder.FenwickTree.sum" l r nFt | otherwise = unsafeSum ft l r --- | Total version of `sum`. Calculates the sum in half-open range \([l, r)\). It returns `Nothing` --- if the interval is invalid. +-- | Total variant of `sum`. Calculates the sum in a half-open interval \([l, r)\). It returns +-- `Nothing` if the interval is invalid. -- -- ==== Complexity -- - \(O(\log n)\) diff --git a/src/AtCoder/Internal/Assert.hs b/src/AtCoder/Internal/Assert.hs --- a/src/AtCoder/Internal/Assert.hs +++ b/src/AtCoder/Internal/Assert.hs @@ -73,7 +73,7 @@ testIndex :: (HasCallStack) => Int -> Int -> Bool testIndex i n = 0 <= i && i < n --- | \(O(1)\) Tests weather \([l, r)\) is a valid interval in \([0, n)\). +-- | \(O(1)\) Tests whether \([l, r)\) is a valid interval in \([0, n)\). -- -- @since 1.0.0.0 {-# INLINE testInterval #-} diff --git a/src/AtCoder/Internal/Barrett.hs b/src/AtCoder/Internal/Barrett.hs --- a/src/AtCoder/Internal/Barrett.hs +++ b/src/AtCoder/Internal/Barrett.hs @@ -1,7 +1,7 @@ {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE RecordWildCards #-} --- | Fast modular multiplication for `Word32` by barrett reduction. +-- | Fast modular multiplication for `Word32` using barrett reduction. -- Reference: https://en.wikipedia.org/wiki/Barrett_reduction -- -- ==== __Example__ @@ -16,10 +16,10 @@ module AtCoder.Internal.Barrett ( -- * Barrett Barrett, - -- * Constructors + -- * Constructor new32, new64, - -- * Accessors + -- * Accessor umod, -- * Barrett reduction mulMod, @@ -29,7 +29,7 @@ import Data.WideWord.Word128 (Word128 (..)) import Data.Word (Word32, Word64) --- | Fast modular multiplication by barrett reduction. +-- | Fast modular multiplication using barrett reduction. -- Reference: https://en.wikipedia.org/wiki/Barrett_reduction -- -- @since 1.0.0.0 @@ -44,14 +44,14 @@ Show ) --- | Creates barret reduction for modulus \(m\) from a `Word32` value. +-- | Creates a `Barrett` for a modulus value \(m\) of type `Word32` value. -- -- @since 1.0.0.0 {-# INLINE new32 #-} new32 :: Word32 -> Barrett new32 m = Barrett m $ maxBound @Word64 `div` (fromIntegral m :: Word64) + 1 --- | Creates barret reduction for modulus \(m\) from a `Word64` value. +-- | Creates a `Barrett` for a modulus value \(m\) of type `Word64` value. -- -- @since 1.0.0.0 {-# INLINE new64 #-} 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 @@ -1,7 +1,6 @@ {-# LANGUAGE RecordWildCards #-} --- | Pushable vector with fixed size capacity. Stack. Internally it tracks the number of elements --- in the vector. +-- | A pushable vector with fixed capacity \(n\). Internally, it tracks the number of elements. -- -- ==== __Example__ -- Create a buffer with capacity @4@: @@ -54,25 +53,21 @@ new, build, - -- * Push/pop - pushBack, - popBack, + -- * Metadata + capacity, + length, + null, - -- * Inspection + -- * Reading back, - - -- * Accessing individual elements read, + + -- * Modifications + pushBack, + popBack, write, modify, modifyM, - - -- * Accesssors - capacity, - length, - null, - - -- * Clearing clear, -- * Conversions @@ -89,7 +84,7 @@ import GHC.Stack (HasCallStack) import Prelude hiding (length, null, read) --- | Pushable vector with fixed size capacity. Stack. +-- | A pushable vector with fixed capacity \(n\). Internally, it tracks the number of elements. -- -- @since 1.0.0.0 data Buffer s a = Buffer @@ -117,30 +112,28 @@ vecB <- VU.thaw xs pure Buffer {..} --- | \(O(1)\) Appends an element to the back. +-- | \(O(1)\) Returns the array size. -- -- @since 1.0.0.0 -{-# INLINE pushBack #-} -pushBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> a -> m () -pushBack Buffer {..} e = do - len <- VGM.read lenB 0 - VGM.write vecB len e - VGM.write lenB 0 (len + 1) +{-# INLINE capacity #-} +capacity :: (VU.Unbox a) => Buffer s a -> Int +capacity = VUM.length . vecB --- | \(O(1)\) Removes the last element from the buffer and returns it, or `Nothing` if it is empty. +-- | \(O(1)\) Returns the number of elements in the buffer. -- -- @since 1.0.0.0 -{-# INLINE popBack #-} -popBack :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m (Maybe a) -popBack Buffer {..} = do - len <- VGM.read lenB 0 - if len == 0 - then pure Nothing - else do - x <- VGM.read vecB (len - 1) - VGM.write lenB 0 (len - 1) - pure $ Just x +{-# INLINE length #-} +length :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m Int +length Buffer {..} = do + VGM.read lenB 0 +-- | \(O(1)\) Returns `True` if the buffer is empty. +-- +-- @since 1.0.0.0 +{-# INLINE null #-} +null :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m Bool +null = (<$>) (== 0) . length + -- | \(O(1)\) Returns the last value in the buffer, or `Nothing` if it is empty. -- -- @since 1.0.0.0 @@ -165,8 +158,32 @@ let !_ = ACIA.checkIndex "AtCoder.Internal.Buffer.read" i len VGM.read vecB i +-- | \(O(1)\) Appends an element to the back. +-- +-- @since 1.0.0.0 +{-# INLINE pushBack #-} +pushBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> a -> m () +pushBack Buffer {..} e = do + len <- VGM.read lenB 0 + VGM.write vecB len e + VGM.write lenB 0 (len + 1) + +-- | \(O(1)\) Removes the last element from the buffer and returns it, or `Nothing` if it is empty. +-- +-- @since 1.0.0.0 +{-# INLINE popBack #-} +popBack :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m (Maybe a) +popBack Buffer {..} = do + len <- VGM.read lenB 0 + if len == 0 + then pure Nothing + else do + x <- VGM.read vecB (len - 1) + VGM.write lenB 0 (len - 1) + pure $ Just x + -- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is --- out of range. +-- out of bounds. -- -- @since 1.0.0.0 {-# INLINE write #-} @@ -177,7 +194,7 @@ VGM.write vecB i e -- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is --- out of range. +-- out of bounds. -- -- @since 1.0.0.0 {-# INLINE modify #-} @@ -188,7 +205,7 @@ VGM.modify vecB f i -- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is --- out of range. +-- out of bounds. -- -- @since 1.0.0.0 {-# INLINE modifyM #-} @@ -197,28 +214,6 @@ len <- VGM.read lenB 0 let !_ = ACIA.checkIndex "AtCoder.Internal.Buffer.modifyM" i len VGM.modifyM vecB f i - --- | \(O(1)\) Returns the array size. --- --- @since 1.0.0.0 -{-# INLINE capacity #-} -capacity :: (VU.Unbox a) => Buffer s a -> Int -capacity = VUM.length . vecB - --- | \(O(1)\) Returns the number of elements in the buffer. --- --- @since 1.0.0.0 -{-# INLINE length #-} -length :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m Int -length Buffer {..} = do - VGM.read lenB 0 - --- | \(O(1)\) Returns `True` if the buffer is empty. --- --- @since 1.0.0.0 -{-# INLINE null #-} -null :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m Bool -null = (<$>) (== 0) . length -- | \(O(1)\) Sets the `length` to zero. -- 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 @@ -55,7 +55,7 @@ Show ) --- | \(O(\log m)\) Creates `FftInfo`. +-- | \(O(\log m)\) Creates an `FftInfo`. -- -- @since 1.0.0.0 {-# INLINE newInfo #-} diff --git a/src/AtCoder/Internal/Csr.hs b/src/AtCoder/Internal/Csr.hs --- a/src/AtCoder/Internal/Csr.hs +++ b/src/AtCoder/Internal/Csr.hs @@ -1,9 +1,10 @@ {-# LANGUAGE RecordWildCards #-} --- | Immutable Compresed Sparse Row. +-- | Immutable Compresed Sparse Row. It is re-exported from the @AtCoder.Extra.Graph@ module with +-- additional functionalities. -- -- ==== __Example__ --- Create a `Csr` without edge weights using `build`: +-- Create a `Csr` without edge weights using `build'` and retrieve the edges with `adj`: -- -- >>> import AtCoder.Internal.Csr qualified as C -- >>> let csr = build' 3 $ VU.fromList @(Int, Int) [(0, 1), (0, 2), (0, 3), (1, 2), (2, 3)] @@ -16,7 +17,7 @@ -- >>> csr `C.adj` 2 -- [3] -- --- Create a `Csr` with edge weights with `build` and retrieve edge with `edgeW`: +-- Create a `Csr` with edge weights using `build` and retrieve the edges with `adjW`: -- -- >>> import AtCoder.Internal.Csr qualified as C -- >>> let csr = build 3 $ VU.fromList @(Int, Int, Int) [(0, 1, 101), (0, 2, 102), (0, 3, 103), (1, 2, 112), (2, 3, 123)] @@ -32,10 +33,13 @@ -- @since 1.0.0.0 module AtCoder.Internal.Csr ( -- * Compressed sparse row - Csr, + Csr (..), + -- * Constructor build, build', + build1, + -- * Accessors adj, adjW, @@ -55,8 +59,25 @@ -- -- @since 1.0.0.0 data Csr w = Csr - { startCsr :: !(VU.Vector Int), + { -- | The number of vertices. + -- + -- @since 1.1.0.0 + nCsr :: {-# UNPACK #-} !Int, + -- | The number of edges. + -- + -- @since 1.1.0.0 + mCsr :: {-# UNPACK #-} !Int, + -- | Starting indices. + -- + -- @since 1.1.0.0 + startCsr :: !(VU.Vector Int), + -- | Adjacent vertices. + -- + -- @since 1.1.0.0 adjCsr :: !(VU.Vector Int), + -- | Edge weights. + -- + -- @since 1.1.0.0 wCsr :: !(VU.Vector w) } deriving @@ -66,25 +87,26 @@ Show ) --- | \(O(n + m)\) Creates `Csr`. +-- | \(O(n + m)\) Creates a `Csr`. -- -- @since 1.0.0.0 {-# INLINE build #-} build :: (HasCallStack, VU.Unbox w) => Int -> VU.Vector (Int, Int, w) -> Csr w -build n edges = runST $ do - start <- VUM.replicate (n + 1) (0 :: Int) +build nCsr edges = runST $ do + let mCsr = VU.length edges + start <- VUM.replicate (nCsr + 1) (0 :: Int) let (!froms, !_, !_) = VU.unzip3 edges VU.forM_ froms $ \from -> do VGM.modify start (+ 1) (from + 1) - for_ [1 .. n] $ \i -> do + for_ [1 .. nCsr] $ \i -> do prev <- VGM.read start (i - 1) VGM.modify start (+ prev) i edgeAdj <- VUM.unsafeNew (VU.length edges) edgeW <- VUM.unsafeNew (VU.length edges) - counter <- VUM.unsafeNew n + counter <- VUM.unsafeNew nCsr VUM.unsafeCopy counter $ VUM.init start VU.forM_ edges $ \(!from, !to, !w) -> do c <- VGM.read counter from @@ -97,7 +119,7 @@ wCsr <- VU.unsafeFreeze edgeW pure Csr {..} --- | \(O(n + m)\) Creates `Csr` with no weight. +-- | \(O(n + m)\) Creates a `Csr` with no edge weight. -- -- @since 1.0.0.0 {-# INLINE build' #-} @@ -106,17 +128,26 @@ where (!us, !vs) = VU.unzip edges --- | \(O(1)\) Returns adjacent vertices. +-- | \(O(n + m)\) Creates a `Csr` with @1@ as edge weights. -- +-- @since 1.1.0.0 +{-# INLINE build1 #-} +build1 :: (HasCallStack) => Int -> VU.Vector (Int, Int) -> Csr Int +build1 n edges = build n $ VU.zip3 us vs (VU.replicate (VU.length us) (1 :: Int)) + where + (!us, !vs) = VU.unzip edges + +-- | \(O(1)\) Returns the adjacent vertices. +-- -- @since 1.0.0.0 {-# INLINE adj #-} -adj :: (HasCallStack, VU.Unbox w) => Csr w -> Int -> VU.Vector Int +adj :: (HasCallStack) => Csr w -> Int -> VU.Vector Int adj Csr {..} i = let il = startCsr VG.! i ir = startCsr VG.! (i + 1) in VU.slice il (ir - il) adjCsr --- | \(O(1)\) Returns adjacent vertices with weights. +-- | \(O(1)\) Returns the adjacent vertices with weights. -- -- @since 1.0.0.0 {-# INLINE adjW #-} @@ -126,11 +157,11 @@ ir = startCsr VG.! (i + 1) in VU.zip (VU.slice il (ir - il) adjCsr) (VU.slice il (ir - il) wCsr) --- | \(O(1)\) Returns a vector of @(edgeId, adjacentVertex)@. +-- | \(O(n)\) Returns a vector of @(edgeId, adjacentVertex)@. -- -- @since 1.0.0.0 {-# INLINE eAdj #-} -eAdj :: (HasCallStack, VU.Unbox w) => Csr w -> Int -> VU.Vector (Int, Int) +eAdj :: (HasCallStack) => Csr w -> Int -> VU.Vector (Int, Int) eAdj Csr {..} i = let il = startCsr VG.! i ir = startCsr VG.! (i + 1) diff --git a/src/AtCoder/Internal/GrowVec.hs b/src/AtCoder/Internal/GrowVec.hs --- a/src/AtCoder/Internal/GrowVec.hs +++ b/src/AtCoder/Internal/GrowVec.hs @@ -39,32 +39,32 @@ -- -- @since 1.0.0.0 module AtCoder.Internal.GrowVec - ( -- * Growable vector + ( -- * GrowVec GrowVec (vecGV), - -- * Constructions - - -- ** Initialization + -- * Constructors new, build, - - -- ** Reserving reserve, - -- * Accessing individual elements + -- * Metadata + length, + capacity, + null, + + -- * Readings read, + + -- * Modifications + + -- ** Writing write, - -- * Modifying the buffer + -- ** Push/pop pushBack, popBack, popBack_, - -- * Accessors - length, - capacity, - null, - -- * Conversion freeze, unsafeFreeze, @@ -91,7 +91,7 @@ vecGV :: !(MutVar s (VUM.MVector s a)) } --- | \(O(n)\) Creates `GrowVec` with initial capacity \(n\). +-- | \(O(n)\) Creates a `GrowVec` with initial capacity \(n\). -- -- @since 1.0.0.0 {-# INLINE new #-} @@ -101,7 +101,7 @@ vecGV <- newMutVar =<< VUM.unsafeNew n pure GrowVec {..} --- | \(O(n)\) Creates `GrowVec` with initial values. +-- | \(O(n)\) Creates a `GrowVec` with initial values. -- -- @since 1.0.0.0 {-# INLINE build #-} @@ -122,6 +122,30 @@ newVec <- VUM.unsafeGrow vec (len - VUM.length vec) writeMutVar vecGV newVec +-- | \(O(1)\) Returns the number of elements in the vector. +-- +-- @since 1.0.0.0 +{-# INLINE length #-} +length :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m Int +length GrowVec {posGV} = do + VGM.unsafeRead posGV 0 + +-- | \(O(1)\) Returns the capacity of the internal the vector. +-- +-- @since 1.0.0.0 +{-# INLINE capacity #-} +capacity :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m Int +capacity GrowVec {vecGV} = do + vec <- readMutVar vecGV + pure $ VUM.length vec + +-- | \(O(1)\) Returns `True` if the vector is empty. +-- +-- @since 1.0.0.0 +{-# INLINE null #-} +null :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m Bool +null = (<$>) (== 0) . length + -- | \(O(1)\) Yields the element at the given position. Will throw an exception if the index is out -- of range. -- @@ -185,7 +209,7 @@ vec <- readMutVar vecGV Just <$> VGM.read vec (pos - 1) --- | \(O(1)\) `popBack` with return value discarded. +-- | \(O(1)\) `popBack` with the return value discarded. -- -- @since 1.0.0.0 {-# INLINE popBack_ #-} @@ -193,30 +217,6 @@ popBack_ GrowVec {..} = do pos <- VGM.unsafeRead posGV 0 VGM.unsafeWrite posGV 0 $ max 0 $ pos - 1 - --- | \(O(1)\) Returns the number of elements in the vector. --- --- @since 1.0.0.0 -{-# INLINE length #-} -length :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m Int -length GrowVec {posGV} = do - VGM.unsafeRead posGV 0 - --- | \(O(1)\) Returns the capacity of the internal the vector. --- --- @since 1.0.0.0 -{-# INLINE capacity #-} -capacity :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m Int -capacity GrowVec {vecGV} = do - vec <- readMutVar vecGV - pure $ VUM.length vec - --- | \(O(1)\) Returns `True` if the vector is empty. --- --- @since 1.0.0.0 -{-# INLINE null #-} -null :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m Bool -null = (<$>) (== 0) . length -- | \(O(n)\) Yields an immutable copy of the mutable vector. -- 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 @@ -64,7 +64,16 @@ -- -- @since 1.0.0.0 {-# INLINE powMod #-} -powMod :: (HasCallStack) => Int -> Int -> Int -> Int +powMod :: + (HasCallStack) => + -- | \(x\) + Int -> + -- | \(n\) + Int -> + -- | \(m\) + Int -> + -- | \(x^n \bmod m\) + Int powMod x n0 m0 | m0 == 1 = 0 | otherwise = fromIntegral $ inner n0 1 $ fromIntegral (x `mod` m0) @@ -81,6 +90,13 @@ -- | M. Forisek and J. Jancina, Fast Primality Testing for Integers That Fit into a Machine Word -- +-- ==== Constraints +-- - \(n < 4759123141 (2^{32} < 4759123141)\), otherwise the return value can lie +-- ([Wikipedia](https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test#Testing_against_small_sets_of_bases)). +-- +-- ==== Complexity +-- - \(O(k \log^3 n)\), \(k = 3\) +-- -- @since 1.0.0.0 {-# INLINE isPrime #-} isPrime :: Int -> Bool @@ -100,7 +116,7 @@ | t == n - 1 || y == 1 || y == n - 1 = not $ y /= n - 1 && even t | otherwise = inner (t .<<. 1) (y * y `mod` n) --- | Returns \((g, x)\) such that \(g = \gcd(a, b), \mathrm{xa} = g(\bmod b), 0 \le x \le b/g\). +-- | Returns \((g, x)\) such that \(g = \gcd(a, b), \mathrm{xa} \equiv g \pmod b, 0 \le x \le b/g\). -- -- ==== Constraints -- - \(1 \le b\) (not asserted) @@ -128,7 +144,7 @@ !m0' = m0 - m1 * u in inner t s' m1 m0' --- | Returns primitive root. +-- | Returns the primitive root of the given `Int`. -- -- @since 1.0.0.0 {-# INLINE primitiveRoot #-} diff --git a/src/AtCoder/Internal/MinHeap.hs b/src/AtCoder/Internal/MinHeap.hs --- a/src/AtCoder/Internal/MinHeap.hs +++ b/src/AtCoder/Internal/MinHeap.hs @@ -34,22 +34,22 @@ ( -- * Heap Heap, - -- * Constructor + -- * Constructors new, - -- * Accessors + -- * Metadata capacity, length, null, - -- * Accessors + -- * Reset clear, - -- * Modifying the heap + -- * Push/pop/peek push, - peek, pop, pop_, + peek, ) where @@ -81,7 +81,7 @@ dataBH :: !(VUM.MVector s a) } --- | \(O(n)\) Creates `Heap` with capacity \(n\). +-- | \(O(n)\) Creates a `Heap` with capacity \(n\). -- -- @since 1.0.0.0 {-# INLINE new #-} @@ -136,17 +136,6 @@ siftUp iParent siftUp i0 --- | \(O(1)\) Returns the smallest value in the heap, or `Nothing` if it is empty. --- --- @since 1.0.0.0 -{-# INLINE peek #-} -peek :: (VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m (Maybe a) -peek heap = do - isNull <- null heap - if isNull - then pure Nothing - else Just <$> VGM.read (dataBH heap) 0 - -- | \(O(\log n)\) Removes the last element from the heap and returns it, or `Nothing` if it is -- empty. -- @@ -189,7 +178,7 @@ siftDown 0 pure $ Just root --- | \(O(\log n)\) `pop` with return value discarded. +-- | \(O(\log n)\) `pop` with the return value discarded. -- -- @since 1.0.0.0 {-# INLINE pop_ #-} @@ -197,3 +186,14 @@ pop_ heap = do _ <- pop heap pure () + +-- | \(O(1)\) Returns the smallest value in the heap, or `Nothing` if it is empty. +-- +-- @since 1.0.0.0 +{-# INLINE peek #-} +peek :: (VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m (Maybe a) +peek heap = do + isNull <- null heap + if isNull + then pure Nothing + else Just <$> VGM.read (dataBH heap) 0 diff --git a/src/AtCoder/Internal/Queue.hs b/src/AtCoder/Internal/Queue.hs --- a/src/AtCoder/Internal/Queue.hs +++ b/src/AtCoder/Internal/Queue.hs @@ -1,6 +1,6 @@ {-# LANGUAGE RecordWildCards #-} --- | Fixed-sized queue. Internally it has \(l, r\) pair of valid element bounds. +-- | Fixed-sized queue. Internally it has an \([l, r)\) pair of valid element bounds. -- -- ==== __Example__ -- >>> import AtCoder.Internal.Queue qualified as Q @@ -43,18 +43,20 @@ -- * Constructor new, - -- * Modifying the queue + -- * Metadata + capacity, + length, + null, + + -- * Modifications + + -- ** Push/pop pushBack, pushFront, popFront, popFront_, - -- * Accessors - capacity, - length, - null, - - -- * Clearing + -- ** Reset clear, -- * Conversions @@ -70,7 +72,7 @@ import GHC.Stack (HasCallStack) import Prelude hiding (length, null) --- | Fixed-sized queue. Internally it has \([l, r)\) pair of valid element bounds. +-- | Fixed-sized queue. Internally it has an \([l, r)\) pair of valid element bounds. -- -- @since 1.0.0.0 data Queue s a = Queue @@ -79,7 +81,7 @@ vecQ :: !(VUM.MVector s a) } --- | \(O(n)\) Creates `Queue` with capacity \(n\). +-- | \(O(n)\) Creates a `Queue` with capacity \(n\). -- -- @since 1.0.0.0 {-# INLINE new #-} @@ -89,6 +91,30 @@ vecQ <- VUM.unsafeNew n pure Queue {..} +-- | \(O(1)\) Returns the array size. +-- +-- @since 1.0.0.0 +{-# INLINE capacity #-} +capacity :: (VU.Unbox a) => Queue s a -> Int +capacity = VUM.length . vecQ + +-- | \(O(1)\) Returns the number of elements in the queue. +-- +-- @since 1.0.0.0 +{-# INLINE length #-} +length :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m Int +length Queue {..} = do + l <- VGM.unsafeRead posQ 0 + r <- VGM.unsafeRead posQ 1 + pure $ r - l + +-- | \(O(1)\) Returns `True` if the buffer is empty. +-- +-- @since 1.0.0.0 +{-# INLINE null #-} +null :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m Bool +null = (<$>) (== 0) . length + -- | \(O(1)\) Appends an element to the back. Will throw an exception if the index is out of range. -- -- @since 1.0.0.0 @@ -136,7 +162,7 @@ VGM.unsafeWrite posQ 0 (l + 1) pure $ Just x --- | \(O(1)\) `popFront` with return value discarded. +-- | \(O(1)\) `popFront` with the return value discarded. -- -- @since 1.0.0.0 {-# INLINE popFront_ #-} @@ -144,30 +170,6 @@ popFront_ que = do _ <- popFront que pure () - --- | \(O(1)\) Returns the array size. --- --- @since 1.0.0.0 -{-# INLINE capacity #-} -capacity :: (VU.Unbox a) => Queue s a -> Int -capacity = VUM.length . vecQ - --- | \(O(1)\) Returns the number of elements in the queue. --- --- @since 1.0.0.0 -{-# INLINE length #-} -length :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m Int -length Queue {..} = do - l <- VGM.unsafeRead posQ 0 - r <- VGM.unsafeRead posQ 1 - pure $ r - l - --- | \(O(1)\) Returns `True` if the buffer is empty. --- --- @since 1.0.0.0 -{-# INLINE null #-} -null :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m Bool -null = (<$>) (== 0) . length -- | \(O(1)\) Sets the `length` to zero. -- 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 @@ -8,15 +8,19 @@ ( -- * Internal SCC SccGraph (nScc), - -- * Constructor + -- * Constructors new, - -- * Modifying the graph + -- * Adding edges addEdge, -- * SCC calculation sccIds, scc, + + -- ** (Extra API) CSR API + sccIdsCsr, + sccCsr ) where @@ -25,6 +29,7 @@ import Control.Monad (unless, when) import Control.Monad.Fix (fix) import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (runST) import Data.Foldable (for_) import Data.Maybe (fromJust) import Data.Vector qualified as V @@ -44,7 +49,7 @@ edgesScc :: !(ACIGV.GrowVec s (Int, Int)) } --- | \(O(n)\) Creates `SccGraph` of \(n\) vertices. +-- | \(O(n)\) Creates a `SccGraph` of \(n\) vertices. -- -- @since 1.0.0.0 {-# INLINE new #-} @@ -67,19 +72,47 @@ {-# INLINE sccIds #-} sccIds :: (PrimMonad m) => SccGraph (PrimState m) -> m (Int, VU.Vector Int) sccIds SccGraph {..} = do + csr <- ACICSR.build' nScc <$> ACIGV.unsafeFreeze edgesScc + pure $ sccIdsCsr csr + +-- | \(O(n + m)\) Returns the strongly connected components. +-- +-- @since 1.0.0.0 +{-# INLINE scc #-} +scc :: (PrimMonad m) => SccGraph (PrimState m) -> m (V.Vector (VU.Vector Int)) +scc g = do + (!groupNum, !ids) <- sccIds g + let counts = VU.create $ do + vec <- VUM.replicate groupNum (0 :: Int) + VU.forM_ ids $ \x -> do + VGM.modify vec (+ 1) x + pure vec + groups <- V.mapM VUM.unsafeNew $ VU.convert counts + is <- VUM.replicate groupNum (0 :: Int) + VU.iforM_ ids $ \v sccId -> do + i <- VGM.read is sccId + VGM.write is sccId $ i + 1 + VGM.write (groups VG.! sccId) i v + V.mapM VU.unsafeFreeze groups + +-- | \(O(n + m)\) API) Returns a pair of @(# of scc, scc id)@. +-- +-- @since 1.1.0.0 +{-# INLINE 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 - g <- ACICSR.build' nScc <$> ACIGV.unsafeFreeze edgesScc -- next SCC ID groupNum <- VUM.replicate 1 (0 :: Int) -- stack of vertices - visited <- ACIGV.new nScc + visited <- ACIGV.new nCsr -- vertex -> low-link: the smallest index of any node on the stack known to be reachable from -- v through v's DFS subtree, including v itself. - low <- VUM.replicate nScc (0 :: Int) + low <- VUM.replicate nCsr (0 :: Int) -- vertex -> order of the visit (0, 1, ..) - ord <- VUM.replicate nScc (-1 :: Int) + ord <- VUM.replicate nCsr (-1 :: Int) -- vertex -> scc id - ids <- VUM.replicate nScc (0 :: Int) + ids <- VUM.replicate nCsr (0 :: Int) let dfs v ord0 = do VGM.write low v ord0 @@ -112,7 +145,7 @@ sccId <- VGM.unsafeRead groupNum 0 fix $ \loop -> do u <- fromJust <$> ACIGV.popBack visited - VGM.write ord u nScc + VGM.write ord u nCsr VGM.write ids u sccId unless (u == v) loop VGM.unsafeWrite groupNum 0 $ sccId + 1 @@ -126,12 +159,12 @@ else pure curOrd ) (0 :: Int) - (VU.generate nScc id) + (VU.generate nCsr id) num <- VGM.unsafeRead groupNum 0 -- The SCCs are reverse topologically sorted, e.g., [0, 1] <- [2] <- [3] -- Now reverse the SCC IDs so that they will be topologically sorted: [3] -> [2] -> [0, 1] - for_ [0 .. nScc - 1] $ \i -> do + for_ [0 .. nCsr - 1] $ \i -> do VGM.modify ids ((num - 1) -) i ids' <- VU.unsafeFreeze ids @@ -139,16 +172,10 @@ -- | \(O(n + m)\) Returns the strongly connected components. -- --- @since 1.0.0.0 -{-# INLINE scc #-} -scc :: (PrimMonad m) => SccGraph (PrimState m) -> m (V.Vector (VU.Vector Int)) -scc g = do - (!groupNum, !ids) <- sccIds g - let counts = VU.create $ do - vec <- VUM.replicate groupNum (0 :: Int) - VU.forM_ ids $ \x -> do - VGM.modify vec (+ 1) x - pure vec +-- @since 1.1.0.0 +{-# INLINE sccCsr #-} +sccCsr :: ACICSR.Csr w -> V.Vector (VU.Vector Int) +sccCsr g = runST $ do groups <- V.mapM VUM.unsafeNew $ VU.convert counts is <- VUM.replicate groupNum (0 :: Int) VU.iforM_ ids $ \v sccId -> do @@ -156,3 +183,10 @@ VGM.write is sccId $ i + 1 VGM.write (groups VG.! sccId) i v V.mapM VU.unsafeFreeze groups + where + (!groupNum, !ids) = sccIdsCsr g + counts = VU.create $ do + vec <- VUM.replicate groupNum (0 :: Int) + VU.forM_ ids $ \x -> do + VGM.modify vec (+ 1) x + pure vec 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 @@ -20,7 +20,7 @@ import Data.Vector.Unboxed.Mutable qualified as VUM import GHC.Stack (HasCallStack) --- TODO: remove `HasCallStack` when we're 100% sure the input is guarded +-- TODO: remove `HasCallStack`? -- | \(O(n^2)\) Internal implementation of suffix array creation (naive). -- @@ -88,7 +88,18 @@ -- -- @since 1.0.0.0 {-# INLINE saIsImpl #-} -saIsImpl :: (HasCallStack) => Int -> Int -> VU.Vector Int -> Int -> VU.Vector Int +saIsImpl :: + (HasCallStack) => + -- | naive threshould + Int -> + -- | doubling threshould + Int -> + -- | string + VU.Vector Int -> + -- | upper bounds + Int -> + -- | suffix array + VU.Vector Int saIsImpl naiveThreshold doublingThreshold s upper = VU.create $ do let n = VU.length s let !ls = VU.create $ do @@ -247,7 +258,14 @@ -- -- @since 1.0.0.0 {-# INLINE saIs #-} -saIs :: (HasCallStack) => VU.Vector Int -> Int -> VU.Vector Int +saIs :: + (HasCallStack) => + -- | string + VU.Vector Int -> + -- | upper bounds + Int -> + -- | suffix array + VU.Vector Int saIs = saIsManual 10 40 -- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting). @@ -259,7 +277,18 @@ -- -- @since 1.0.0.0 {-# INLINE saIsManual #-} -saIsManual :: (HasCallStack) => Int -> Int -> VU.Vector Int -> Int -> VU.Vector Int +saIsManual :: + (HasCallStack) => + -- | naive threshold + Int -> + -- | doubling threshold + Int -> + -- | string + VU.Vector Int -> + -- | upper bounds + Int -> + -- | suffix array + VU.Vector Int saIsManual naiveThreshold doublingThreshold s upper | n == 0 = VU.empty | n == 1 = VU.singleton 0 diff --git a/src/AtCoder/LazySegTree.hs b/src/AtCoder/LazySegTree.hs --- a/src/AtCoder/LazySegTree.hs +++ b/src/AtCoder/LazySegTree.hs @@ -1,6 +1,6 @@ {-# LANGUAGE RecordWildCards #-} --- | Lazily propagted segment tree. It is the data structure for the pair of a [monoid](https://en.wikipedia.org/wiki/Monoid) +-- | A lazily propagted segment tree. It is the data structure for the pair of a [monoid](https://en.wikipedia.org/wiki/Monoid) -- \((S, \cdot: S \times S \to S, e \in S)\) and a set \(F\) of \(S \to S\) mappings that satisfies -- the following properties. -- @@ -51,7 +51,7 @@ -- >>> LST.allProd seg -- Sum {getSum = 24} -- --- Run binary search in \(O(\log n\) time complexity: +-- Run binary search: -- -- >>> LST.maxRight seg 0 (<= (Sum 10)) -- sum [0, 2) = 7 <= 10 -- 2 @@ -59,8 +59,7 @@ -- >>> LST.minLeft seg 4 (<= (Sum 10)) -- sum [3, 4) = 10 <= 10 -- 3 -- --- Inspect all the values in \(O(n \log n)\) with `freeze` or `unsafeFreeze`. Note that they --- propagete all the applied actions: +-- Inspect all the values in \(O(n \log n)\) with `freeze` or `unsafeFreeze`: -- -- >>> VU.map getSum <$> LST.freeze seg -- [2,5,7,10] @@ -69,14 +68,14 @@ -- -- - `prod` returns \(a_l \cdot a_{l + 1} \cdot .. \cdot a_{r - 1}\). If you need \(a_{r - 1} \cdot a_{r - 2} \cdot .. \cdot a_{l}\), -- wrap your monoid in `Data.Monoid.Dual`. --- - If you ever need to store boxed types to `LazySegTree`, wrap it in 'vector:Data.Vector.Unboxed.DoNotUnboxStrict' +-- - If you ever need to store boxed types to `LazySegTree`, wrap it in @Data.Vector.Unboxed.DoNotUnboxStrict@ -- or the like. -- -- ==== Major changes from the original @ac-library@ -- - The API is based on `Monoid` and `SegAct`, not the functions @op@, @e@, @mapping@, -- @composition@ and @id@. -- - @get@ and @set@ are renamed to `read` and `write`. --- - `modify`, `modifyM`, `freeze` and `unsafeFreeze` are added. +-- - `modify`, `modifyM`, `exchange`, `freeze` and `unsafeFreeze` are added. -- -- @since 1.0.0.0 module AtCoder.LazySegTree @@ -88,10 +87,11 @@ new, build, - -- * Accessing individual elements + -- * Accessing elements write, modify, modifyM, + exchange, read, -- * Products @@ -134,17 +134,19 @@ -- | Typeclass reprentation of the `LazySegTree` properties. User can implement either `segAct` or -- `segActWithLength`. -- --- Instances should satisfy the follwing: +-- Instances should satisfy the follwing properties: -- -- [Left monoid action] @'segAct' (f2 '<>' f1) x = 'segAct' f2 ('segAct' f1 x)@ -- [Identity map] @`segAct` `mempty` x = x@ -- [Endomorphism] @'segAct' f (x1 '<>' x2) = ('segAct' f x1) '<>' ('segAct' f x2)@ -- --- If you implement `segActWithLength`, satisfy one more propety: +-- If you implement `SegAct` via `segActWithLength`, satisfy one more propety: -- -- [Linear left monoid action] @'segActWithLength' len f a = 'Data.Semigroup.stimes' len ('segAct' f a) a@. -- --- Note that in `SegAct` instances, new semigroup values are always given from the left: @new '<>' old@. +-- ==== Invariant +-- In `SegAct` instances, new semigroup values are always given from the left: @new '<>' old@. The +-- order is important for non-commutative monoid implementations. -- -- ==== __Example instance__ -- Take `AtCoder.Extra.Monoid.Affine1` as an example of type \(F\). @@ -287,7 +289,7 @@ segActWithLength :: Int -> f -> a -> a segActWithLength _ = segAct --- | Lazy segment tree defined around `SegAct`. +-- | A lazily propagated segment tree defined around `SegAct`. -- -- @since 1.0.0.0 data LazySegTree s f a = LazySegTree @@ -408,6 +410,27 @@ for_ [1 .. logLst] $ \i -> do update self $ p' .>>. i +-- | (Extra API) Sets \(p\)-th value of the array to \(x\) and returns the old value. +-- +-- ==== Constraints +-- - \(0 \leq p \lt n\) +-- +-- ==== Complexity +-- - \(O(\log n)\) +-- +-- @since 1.1.0.0 +{-# INLINE exchange #-} +exchange :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> a -> m a +exchange self@LazySegTree {..} p x = do + let !_ = ACIA.checkIndex "AtCoder.LazySegTree.exchange" p nLst + let p' = p + sizeLst + for_ [logLst, logLst - 1 .. 1] $ \i -> do + push self $ p' .>>. i + res <- VGM.exchange dLst p' x + for_ [1 .. logLst] $ \i -> do + update self $ p' .>>. i + pure res + -- | Returns \(p\)-th value of the array. -- -- ==== Constraints @@ -443,8 +466,8 @@ | l0 == r0 = pure mempty | otherwise = unsafeProd self l0 r0 --- | Total version of `prod`. Returns the product of \([a[l], ..., a[r - 1]]\), assuming the --- properties of the monoid. It returns `Just` `mempty` if \(l = r\). It returns `Nothing` if the +-- | Total variant of `prod`. Returns the product of \([a[l], ..., a[r - 1]]\), assuming the +-- properties of the monoid. Returns `Just` `mempty` if \(l = r\). It returns `Nothing` if the -- interval is invalid. -- -- ==== Complexity @@ -494,7 +517,7 @@ allProd :: (PrimMonad m, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> m a allProd LazySegTree {..} = VGM.read dLst 1 --- | Applies @segAct f@ to an index @p@. +-- | Applies @segAct f@ to an index \(p\). -- -- ==== Constraints -- - \(0 \leq p \lt n\) @@ -517,7 +540,7 @@ for_ [1 .. logLst] $ \i -> do update self $ p' .>>. i --- | Applies @segAct f@ to an interval @[l, r)@. +-- | Applies @segAct f@ to an interval \([l, r)\). -- -- ==== Constraints -- - \(0 \leq l \leq r \leq n\) @@ -577,7 +600,7 @@ minLeft :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> (a -> Bool) -> m Int minLeft seg r0 g = minLeftM seg r0 (pure . g) --- | Monadic version of `minLeft`. +-- | Monadic variant of `minLeft`. -- -- ==== Constraints -- @@ -652,7 +675,7 @@ maxRight :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> (a -> Bool) -> m Int maxRight seg l0 g = maxRightM seg l0 (pure . g) --- | Monadic version of `maxRight`. +-- | Monadic variant of `maxRight`. -- -- ==== Constraints -- diff --git a/src/AtCoder/Math.hs b/src/AtCoder/Math.hs --- a/src/AtCoder/Math.hs +++ b/src/AtCoder/Math.hs @@ -5,6 +5,7 @@ -- @since 1.0.0.0 module AtCoder.Math ( -- * Modulus operations + -- These functions are internally used for `AtCoder.ModInt`. powMod, invMod, @@ -42,7 +43,14 @@ -- -- @since 1.0.0.0 {-# INLINE invMod #-} -invMod :: (HasCallStack) => Int -> Int -> Int +invMod :: + (HasCallStack) => + -- | \(x\) + Int -> + -- | \(m\) + Int -> + -- | \(x^{-1} \bmod m\) + Int invMod x m = let !_ = ACIA.runtimeAssert (1 <= m) $ "AtCoder.Math.invMod: given invalid `m` less than 1: " ++ show m (!z1, !z2) = ACIM.invGcd (fromIntegral x) (fromIntegral m) @@ -66,7 +74,7 @@ -- ==== Complexity -- - \(O(n \log{\mathrm{lcm}(m[i])})\) -- --- ==== Example +-- ==== __Example__ -- `crt` calculates \(y\) such that \(y \equiv r_i \pmod m_i, \forall i \in \lbrace 0,1,\cdots, n - 1 \rbrace\): -- -- >>> import Data.Vector.Unboxed qualified as VU @@ -127,7 +135,7 @@ -- ==== Complexity -- - \(O(\log m)\) -- --- ==== Example +-- ==== __Example__ -- `floorSum` calculates the number of points surrounded by a line -- \(y = \frac {a \times x + b} {m} \) and \(x, y\) axes in \(O(\log m)\) time: -- @@ -150,7 +158,18 @@ -- -- @since 1.0.0.0 {-# INLINE floorSum #-} -floorSum :: (HasCallStack) => Int -> Int -> Int -> Int -> Int +floorSum :: + (HasCallStack) => + -- | \(n\) + Int -> + -- | \(m\) + Int -> + -- | \(a\) + Int -> + -- | \(b\) + Int -> + -- | \(\sum\limits_{i = 0}^{n - 1} \left\lfloor \frac{a \times i + b}{m} \right\rfloor\) + Int floorSum n m a b = ACIM.floorSumUnsigned n m a' b' - da - db where !_ = ACIA.runtimeAssert (0 <= n && n < bit 32) $ "AtCoder.Math.floorSum: given invalid `n` (`" ++ show n ++ "`)" diff --git a/src/AtCoder/MaxFlow.hs b/src/AtCoder/MaxFlow.hs --- a/src/AtCoder/MaxFlow.hs +++ b/src/AtCoder/MaxFlow.hs @@ -21,7 +21,7 @@ -- 1 -- -- Get the minimum cut with `minCut`. In this case, removing the second edge makes the minimum cut --- (note that the edge capacity (\(1\)) = max flow): +-- (note that the edge capacity \(1\) = max flow): -- -- >>> MF.minCut g 0 -- returns a Bit vector. `1` (`Bit True`) is on the `s` side. -- [1,1,0] @@ -116,7 +116,18 @@ -- -- @since 1.0.0.0 {-# INLINE addEdge #-} -addEdge :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => MfGraph (PrimState m) cap -> Int -> Int -> cap -> m Int +addEdge :: + (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => + -- | Graph + MfGraph (PrimState m) cap -> + -- | from + Int -> + -- | to + Int -> + -- | cap + cap -> + -- | Edge index + m Int addEdge MfGraph {..} from to cap = do let !_ = ACIA.checkCustom "AtCoder.MaxFlow.addEdge" "`from` vertex" from "the number of vertices" nG let !_ = ACIA.checkCustom "AtCoder.MaxFlow.addEdge" "`to` vertex" to "the number of vertices" nG @@ -131,7 +142,7 @@ ACIGV.pushBack (gG VG.! to) (from, iEdge, 0) pure m --- | `addEdge` with return value discarded. +-- | `addEdge` with the return value discarded. -- -- ==== Constraints -- - \(0 \leq \mathrm{from}, \mathrm{to} \lt n\) @@ -142,7 +153,17 @@ -- -- @since 1.0.0.0 {-# INLINE addEdge_ #-} -addEdge_ :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => MfGraph (PrimState m) cap -> Int -> Int -> cap -> m () +addEdge_ :: + (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => + -- | Graph + MfGraph (PrimState m) cap -> + -- | from + Int -> + -- | to + Int -> + -- | cap + cap -> + m () addEdge_ graph from to cap = do _ <- addEdge graph from to cap pure () @@ -161,11 +182,22 @@ -- -- @since 1.0.0.0 {-# INLINE flow #-} -flow :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => MfGraph (PrimState m) cap -> Int -> Int -> cap -> m cap +flow :: + (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => + -- | Graph + MfGraph (PrimState m) cap -> + -- | Source @s@ + Int -> + -- | Sink @t@ + Int -> + -- | Flow limit + cap -> + -- | Max flow + m cap flow MfGraph {..} s t flowLimit = do let !_ = ACIA.checkCustom "AtCoder.MaxFlow.flow" "`source` vertex" s "the number of vertices" nG let !_ = ACIA.checkCustom "AtCoder.MaxFlow.flow" "`sink` vertex" t "the number of vertices" nG - let !_ = ACIA.runtimeAssert (s /= t) $ "AtCoder.MaxFlow.flow: `source` and `sink` vertex have to be distinct: `" ++ show s ++ "`" + let !_ = ACIA.runtimeAssert (s /= t) $ "AtCoder.MaxFlow.flow: `source` and `sink` vertex must be distinct: `" ++ show s ++ "`" level <- VUM.unsafeNew nG que <- ACIQ.new nG @@ -253,7 +285,16 @@ -- -- @since 1.0.0.0 {-# INLINE maxFlow #-} -maxFlow :: (HasCallStack, PrimMonad m, Num cap, Ord cap, Bounded cap, VU.Unbox cap) => MfGraph (PrimState m) cap -> Int -> Int -> m cap +maxFlow :: + (HasCallStack, PrimMonad m, Num cap, Ord cap, Bounded cap, VU.Unbox cap) => + -- | Graph + MfGraph (PrimState m) cap -> + -- | Source @s@ + Int -> + -- | Sink @t@ + Int -> + -- | Max flow + m cap maxFlow graph s t = flow graph s t maxBound -- | Returns a vector of length \(n\), such that the \(i\)-th element is `True` if and only if there @@ -265,7 +306,14 @@ -- -- @since 1.0.0.0 {-# INLINE minCut #-} -minCut :: (PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => MfGraph (PrimState m) cap -> Int -> m (VU.Vector Bit) +minCut :: + (PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => + -- | Graph + MfGraph (PrimState m) cap -> + -- | Source @s@ + Int -> + -- | Minimum cut + m (VU.Vector Bit) minCut MfGraph {..} s = do visited <- VUM.replicate nG $ Bit False que <- ACIQ.new nG -- we could use a growable queue here @@ -296,7 +344,14 @@ -- -- @since 1.0.0.0 {-# INLINE getEdge #-} -getEdge :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => MfGraph (PrimState m) cap -> Int -> m (Int, Int, cap, cap) +getEdge :: + (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => + -- | Graph + MfGraph (PrimState m) cap -> + -- | Vertex + Int -> + -- | Tuple of @(from, to, cap, flow)@ + m (Int, Int, cap, cap) getEdge MfGraph {..} i = do m <- ACIGV.length posG let !_ = ACIA.checkEdge "AtCoder.MaxFlow.getEdge" i m @@ -313,13 +368,18 @@ -- -- @since 1.0.0.0 {-# INLINE edges #-} -edges :: (PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => MfGraph (PrimState m) cap -> m (VU.Vector (Int, Int, cap, cap)) +edges :: + (PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => + -- | Graph + MfGraph (PrimState m) cap -> + -- | Vector of @(from, to, cap, flow)@ + m (VU.Vector (Int, Int, cap, cap)) edges g@MfGraph {posG} = do len <- ACIGV.length posG VU.generateM len (getEdge g) -- | \(O(1)\) Changes the capacity and the flow amount of the $i$-th edge to @newCap@ and --- @newFlow@, respectively. It doesn't change the capacity or the flow amount of other edges. +-- @newFlow@, respectively. It oes not change the capacity or the flow amount of other edges. -- -- ==== Constraints -- - \(0 \leq \mathrm{newflow} \leq \mathrm{newcap}\) @@ -329,7 +389,17 @@ -- -- @since 1.0.0.0 {-# INLINE changeEdge #-} -changeEdge :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => MfGraph (PrimState m) cap -> Int -> cap -> cap -> m () +changeEdge :: + (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => + -- | Graph + MfGraph (PrimState m) cap -> + -- | Edge index + Int -> + -- | New capacity + cap -> + -- | New flow + cap -> + m () changeEdge MfGraph {..} i newCap newFlow = do m <- ACIGV.length posG let !_ = ACIA.checkEdge "AtCoder.MaxFlow.changeEdge" i m diff --git a/src/AtCoder/MinCostFlow.hs b/src/AtCoder/MinCostFlow.hs --- a/src/AtCoder/MinCostFlow.hs +++ b/src/AtCoder/MinCostFlow.hs @@ -110,11 +110,17 @@ {-# INLINE addEdge #-} addEdge :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) => + -- | Graph McfGraph (PrimState m) cap cost -> + -- | from Int -> + -- | to Int -> + -- | capacity cap -> + -- | cost cost -> + -- | Edge index m Int addEdge McfGraph {..} from to cap cost = do let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.addEdge" "`from` vertex" from "the number of vertices" nG @@ -125,7 +131,7 @@ ACIGV.pushBack edgesG (from, to, cap, 0, cost) pure m --- | `addEdge` with return value discarded. +-- | `addEdge` with the return value discarded. -- -- ==== Constraints -- - \(0 \leq \mathrm{from}, \mathrm{to} \lt n\) @@ -138,10 +144,15 @@ {-# INLINE addEdge_ #-} addEdge_ :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) => + -- | Graph McfGraph (PrimState m) cap cost -> + -- | from Int -> + -- | to Int -> + -- | capacity cap -> + -- | cost cost -> m () addEdge_ graph from to cap cost = do @@ -161,10 +172,15 @@ {-# INLINE flow #-} flow :: (HasCallStack, PrimMonad m, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) => + -- | Graph McfGraph (PrimState m) cap cost -> + -- | Fource @s@ Int -> + -- | Sink @t@ Int -> + -- | Flow limit cap -> + -- | Tuple of @(cap, cost@) m (cap, cost) flow graph s t flowLimit = do res <- slope graph s t flowLimit @@ -182,9 +198,13 @@ {-# INLINE maxFlow #-} maxFlow :: (HasCallStack, PrimMonad m, Integral cap, Ord cap, Bounded cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) => + -- | Graph McfGraph (PrimState m) cap cost -> + -- | Source @s@ Int -> + -- | Sink @t@ Int -> + -- | Tuple of @(cap, cost@) m (cap, cost) maxFlow graph s t = do res <- slope graph s t maxBound @@ -216,15 +236,20 @@ {-# INLINE slope #-} slope :: (HasCallStack, PrimMonad m, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) => + -- | Graph McfGraph (PrimState m) cap cost -> + -- | Source @s@ Int -> + -- | Sink @t@ Int -> + -- | Flow limit cap -> + -- | Vector of @(cap, cost)@ m (VU.Vector (cap, cost)) slope McfGraph {..} s t flowLimit = do let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.slope" "`source` vertex" s "the number of vertices" nG let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.slope" "`sink` vertex" t "the number of vertices" nG - let !_ = ACIA.runtimeAssert (s /= t) "AtCoder.MinCostFlow.slope: `source` and `sink` vertex have to be distict" + let !_ = ACIA.runtimeAssert (s /= t) "AtCoder.MinCostFlow.slope: `source` and `sink` vertex must be distict" edges_@(VU.V_5 _ _ _ caps _ _) <- ACIGV.unsafeFreeze edgesG (!edgeIdx, !g) <- ACIMCSR.build nG edges_ @@ -292,8 +317,8 @@ cap <- VGM.read capCsr $ start + di unless (cap == 0) $ do - -- \|-dual[e.to] + dual[v]| <= (n-1)C - -- cost <= C - -(n-1)C + 0 = nC + -- - |-dual[e.to] + dual[v]| <= (n-1)C + -- - cost <= C - -(n-1)C + 0 = nC cost' <- do dualTo <- VGM.read duals to pure $! cost - dualTo + dualV @@ -368,8 +393,11 @@ {-# INLINE getEdge #-} getEdge :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) => + -- | Graph McfGraph (PrimState m) cap cost -> + -- | Edge index Int -> + -- | Tuple of @(from, to, cap, flow, cost)@ m (Int, Int, cap, cap, cost) getEdge McfGraph {..} i = do m <- ACIGV.length edgesG @@ -386,7 +414,9 @@ {-# INLINE edges #-} edges :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) => + -- | Graph McfGraph (PrimState m) cap cost -> + -- | Vector of @(from, to, cap, flow, cost)@ m (VU.Vector (Int, Int, cap, cap, cost)) edges McfGraph {..} = do ACIGV.freeze edgesG @@ -401,7 +431,9 @@ {-# INLINE unsafeEdges #-} unsafeEdges :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) => + -- | Graph McfGraph (PrimState m) cap cost -> + -- | Vector of @(from, to, cap, flow, cost)@ m (VU.Vector (Int, Int, cap, cap, cost)) unsafeEdges McfGraph {..} = do ACIGV.unsafeFreeze edgesG diff --git a/src/AtCoder/ModInt.hs b/src/AtCoder/ModInt.hs --- a/src/AtCoder/ModInt.hs +++ b/src/AtCoder/ModInt.hs @@ -1,11 +1,11 @@ {-# LANGUAGE MagicHash #-} {-# LANGUAGE TypeFamilies #-} --- | It is the struct that treats the modular arithmetic. All the remaining parts of AC Library +-- | It is the structure that treats the modular arithmetic. All the remaining parts of AC Library -- works without modint, so you don't necessarily read this to use the remaining parts. -- --- For most of the problems, it is sufficient to use `ModInt998244353`, `ModInt1000000007`, which --- can be used as follows. +-- ==== __Example__ +-- It is often convenient to define a type alias of `ModInt` for a specific modulus value: -- -- >>> import AtCoder.ModInt qualified as M -- >>> type Mint = M.ModInt998244353 @@ -13,7 +13,11 @@ -- >>> modInt 1000000000 -- 1755647 -- +-- >>> modInt 1000000000 / modInt 3 +-- 666081451 +-- -- ==== Major changes from the original @ac-library@ +-- - @StaticModInt@ is renamed to `ModInt`. -- - @DynamicModInt@ is removed. -- -- @since 1.0.0.0 @@ -82,14 +86,14 @@ -- @since 1.0.0.0 isPrimeModulus :: Proxy# a -> Bool - -- | Returns the primitive root of the modulus value. Note that the default implementation is - -- slow. + -- | Returns the primitive root of the modulus value. Note that the default implementation is slow + -- and the value should be hard-coded. -- -- @since 1.0.0.0 {-# INLINE primitiveRootModulus #-} primitiveRootModulus :: Proxy# a -> Int -- we could use `AllowAmbigousTypes` or `Tagged` newtype, but `Proxy#` wasn't so slow. - -- not sure about `x^n` case though.. + -- not sure about the case of `x^n` though.. primitiveRootModulus _ = ACIM.primitiveRoot $ fromIntegral (natVal' (proxy# @a)) -- | \(2^{24} - 1\). @@ -119,7 +123,7 @@ {-# INLINE primitiveRootModulus #-} primitiveRootModulus _ = 11 --- | \(119 \times 2^{23} + 1\). It is often used in contest problems +-- | \(119 \times 2^{23} + 1\). It is often used in contest problems. -- -- @since 1.0.0.0 instance Modulus 998244353 where @@ -146,17 +150,13 @@ {-# INLINE primitiveRootModulus #-} primitiveRootModulus _ = 7 --- | `ModInt` with modulus value @998244353@. --- --- @since 1.0.0.0 +-- | @since 1.0.0.0 type ModInt998244353 = ModInt 998244353 --- | `ModInt` with modulus value @1000000007@. --- --- @since 1.0.0.0 +-- | @since 1.0.0.0 type ModInt1000000007 = ModInt 1000000007 --- | Retrieves `Int` from `KnownNat`. +-- | Retrieves the `Int` value from a `KnownNat`. -- -- >>> import Data.Proxy (Proxy(..)) -- >>> modVal (Proxy @42) @@ -167,7 +167,7 @@ modVal :: forall a. (KnownNat a) => Proxy a -> Int modVal p = fromIntegral $ natVal p --- | Retrieves `Int` from `KnownNat`. +-- | Retrieves the `Int` value from a `KnownNat`. -- -- >>> :set -XMagicHash -- >>> import GHC.Exts (proxy#) @@ -179,28 +179,28 @@ modVal# :: forall a. (KnownNat a) => Proxy# a -> Int modVal# p = fromIntegral $ natVal' p --- | Creates `ModInt` from an `Int` value taking mod. +-- | Creates a `ModInt` from an `Int` value taking the mod. -- -- @since 1.0.0.0 {-# INLINE new #-} new :: forall a. (KnownNat a) => Int -> ModInt a new v = ModInt . fromIntegral $ v `mod` fromIntegral (natVal' (proxy# @a)) --- | Creates `ModInt` from a `Word32` value taking mod. +-- | Creates a `ModInt` from a `Word32` value taking the mod. -- -- @since 1.0.0.0 {-# INLINE new32 #-} new32 :: forall a. (KnownNat a) => Word32 -> ModInt a new32 v = ModInt $ v `mod` fromIntegral (natVal' (proxy# @a)) --- | Creates `ModInt` from a `Word64` value taking mod. +-- | Creates a `ModInt` from a `Word64` value taking the mod. -- -- @since 1.0.0.0 {-# INLINE new64 #-} new64 :: forall a. (KnownNat a) => Word64 -> ModInt a new64 v = ModInt . fromIntegral $ v `mod` fromIntegral (natVal' (proxy# @a)) --- | Creates `ModInt` without taking mod. It is the function for constant-factor speedup. +-- | Creates `ModInt` without taking the mod. It is the function for constant-factor speedup. -- -- ==== Constraints -- - \(0 \leq x \lt \mathrm{mod}\) (not asserted at runtime) @@ -210,20 +210,23 @@ unsafeNew :: (KnownNat a) => Word32 -> ModInt a unsafeNew = ModInt --- | `Word32` value that treats the modula arithmetic. -newtype ModInt a = ModInt {unModInt :: Word32} +-- | `Word32` value that treats the modular arithmetic. +newtype ModInt a = ModInt + { -- | @since 1.0.0.0 + unModInt :: Word32 + } deriving - ( -- @since 1.0.0.0 + ( -- | @since 1.0.0.0 P.Prim ) deriving newtype - ( -- @since 1.0.0.0 + ( -- | @since 1.0.0.0 Eq, - -- @since 1.0.0.0 + -- | @since 1.0.0.0 Ord, - -- @since 1.0.0.0 + -- | @since 1.0.0.0 Read, - -- @since 1.0.0.0 + -- | @since 1.0.0.0 Show ) @@ -322,10 +325,10 @@ !_ = ACIA.runtimeAssert (eg1 == 1) "AtCoder.ModInt.inv: `x^(-1) mod m` cannot be calculated when `gcd x modulus /= 1`" in fromIntegral eg2 --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 deriving newtype instance (KnownNat p) => Real (ModInt p) --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 instance (KnownNat p) => Num (ModInt p) where {-# INLINE (+) #-} (ModInt !x1) + (ModInt !x2) @@ -355,45 +358,45 @@ {-# INLINE fromInteger #-} fromInteger = ModInt . fromInteger . (`mod` fromIntegral (natVal' (proxy# @p))) --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 instance (KnownNat p) => Bounded (ModInt p) where {-# INLINE minBound #-} minBound = ModInt 0 {-# INLINE maxBound #-} maxBound = ModInt $! fromIntegral (natVal' (proxy# @p)) - 1 --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 instance (KnownNat p) => Enum (ModInt p) where {-# INLINE toEnum #-} toEnum = new {-# INLINE fromEnum #-} fromEnum = fromIntegral . unModInt --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 instance (Modulus p) => Integral (ModInt p) where {-# INLINE quotRem #-} quotRem x y = (x / y, x - x / y * y) {-# INLINE toInteger #-} toInteger = coerce (toInteger @Word32) --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 instance (Modulus p) => Fractional (ModInt p) where {-# INLINE recip #-} recip = inv {-# INLINE fromRational #-} fromRational q = fromInteger (numerator q) / fromInteger (denominator q) --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 newtype instance VU.MVector s (ModInt a) = MV_ModInt (VU.MVector s Word32) --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 newtype instance VU.Vector (ModInt a) = V_ModInt (VU.Vector Word32) --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 deriving newtype instance VGM.MVector VU.MVector (ModInt a) --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 deriving newtype instance VG.Vector VU.Vector (ModInt a) --- | -- @since 1.0.0.0 +-- | @since 1.0.0.0 instance VU.Unbox (ModInt a) diff --git a/src/AtCoder/Scc.hs b/src/AtCoder/Scc.hs --- a/src/AtCoder/Scc.hs +++ b/src/AtCoder/Scc.hs @@ -1,16 +1,22 @@ -- | It calculates the strongly connected components of directed graphs. -- --- ==== Example +-- ==== __Example__ +-- Create a `SccGraph`: +-- -- >>> import AtCoder.Scc qualified as Scc -- >>> gr <- Scc.new 4 -- 0 1 2 3 -- >>> Scc.nScc gr -- 4 -- +-- Add edges and get SCC of the graph: +-- -- >>> Scc.addEdge gr 0 1 -- 0 -> 1 2 3 -- >>> Scc.addEdge gr 1 0 -- 0 == 1 2 3 -- >>> Scc.addEdge gr 1 2 -- 0 == 1 -> 2 3 -- >>> Scc.scc gr -- [[3],[0,1],[2]] +-- +-- See also the @scc@ function in @AtCoder.Extra.Graph@ module that computes SCC for a CSR. -- -- @since 1.0.0.0 module AtCoder.Scc (SccGraph, nScc, new, addEdge, scc) where diff --git a/src/AtCoder/SegTree.hs b/src/AtCoder/SegTree.hs --- a/src/AtCoder/SegTree.hs +++ b/src/AtCoder/SegTree.hs @@ -21,6 +21,7 @@ -- Create a `SegTree` of @'Sum' Int@: -- -- >>> import AtCoder.SegTree qualified as ST +-- >>> import Data.Vector.Unboxed qualified as VU -- >>> import Data.Monoid (Sum(..)) -- >>> seg <- ST.new @_ @(Sum Int) 4 -- @@ -57,13 +58,13 @@ -- -- - `prod` returns \(a_l \cdot a_{l + 1} \cdot .. \cdot a_{r - 1}\). If you need \(a_{r - 1} \cdot a_{r - 2} \cdot .. \cdot a_{l}\), -- wrap your monoid in `Data.Monoid.Dual`. --- - If you ever need to store boxed types to `LazySegTree`, wrap it in 'vector:Data.Vector.Unboxed.DoNotUnboxStrict' +-- - If you ever need to store boxed types to `LazySegTree`, wrap it in @Data.Vector.Unboxed.DoNotUnboxStrict@ -- or the like. -- -- ==== Major changes from the original @ac-library@ --- - The implementation is `Monoid` based. +-- - The implementation is `Monoid` based, not function objects. -- - @get@ and @set@ are renamed to `read` and `write`. --- - `modify`, `modifyM`, `freeze` and `unsafeFreeze` are added. +-- - `modify`, `modifyM`, `exchange`, `freeze` and `unsafeFreeze` are added. -- -- @since 1.0.0.0 module AtCoder.SegTree @@ -74,10 +75,11 @@ new, build, - -- * Accessing individual elements + -- * Accessing elements write, modify, modifyM, + exchange, read, -- * Products @@ -145,7 +147,7 @@ new :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Int -> m (SegTree (PrimState m) a) new nSt | nSt >= 0 = build $ VU.replicate nSt mempty - | otherwise = error $ "new: given negative size (`" ++ show nSt ++ "`)" + | otherwise = error $ "AtCoder.SegTree.new: given negative size (`" ++ show nSt ++ "`)" -- | Creates an array with initial values. -- @@ -218,6 +220,25 @@ for_ [1 .. logSt] $ \i -> do update self ((p + sizeSt) .>>. i) +-- | (Extra API) Writes \(p\)-th value of the array to \(x\) and returns the old value. +-- +-- ==== Constraints +-- - \(0 \leq p \lt n\) +-- +-- ==== Complexity +-- - \(O(\log n)\) +-- +-- @since 1.1.0.0 +{-# INLINE exchange #-} +exchange :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> a -> m a +exchange self@SegTree {..} p x = do + let !_ = ACIA.checkIndex "AtCoder.SegTree.exchange" p nSt + ret <- VGM.exchange dSt (p + sizeSt) x + VGM.write dSt (p + sizeSt) x + for_ [1 .. logSt] $ \i -> do + update self ((p + sizeSt) .>>. i) + pure ret + -- | Returns \(p\)-th value of the array. -- -- ==== Constraints @@ -249,8 +270,8 @@ | ACIA.testInterval l0 r0 nSt = unsafeProd self l0 r0 | otherwise = ACIA.errorInterval "AtCoder.SegTree.prod" l0 r0 nSt --- | Total version of `prod`. Returns \(a[l] \cdot ... \cdot a[r - 1]\), assuming the properties of --- the monoid. It returns `Just` `mempty` if \(l = r\). It return `Nothing` if the interval is +-- | Total variant of `prod`. Returns \(a[l] \cdot ... \cdot a[r - 1]\), assuming the properties of +-- the monoid. It returns `Just` `mempty` if \(l = r\). Returns `Nothing` if the interval is -- invalid. -- -- ==== Complexity @@ -315,10 +336,19 @@ -- -- @since 1.0.0.0 {-# INLINE minLeft #-} -minLeft :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> (a -> Bool) -> m Int +minLeft :: + (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => + -- | The segment tree + SegTree (PrimState m) a -> + -- | \(r\) + Int -> + -- | \(p\): user prediate + (a -> Bool) -> + -- | \(l\): \(p\) holds for \([l, r)\) + m Int minLeft seg r0 f = minLeftM seg r0 (pure . f) --- | Monadic version of `minLeft`. +-- | Monadic variant of `minLeft`. -- -- ==== Constraints -- @@ -332,7 +362,16 @@ -- -- @since 1.0.0.0 {-# INLINE minLeftM #-} -minLeftM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int +minLeftM :: + (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => + -- | The segment tree + SegTree (PrimState m) a -> + -- | \(r\) + Int -> + -- | \(p\): user prediate + (a -> m Bool) -> + -- | \(l\): \(p\) holds for \([l, r)\) + m Int minLeftM SegTree {..} r0 f = do b <- f mempty let !_ = ACIA.runtimeAssert b "AtCoder.SegTree.minLeftM: `f empty` returned `False`" @@ -385,10 +424,19 @@ -- -- @since 1.0.0.0 {-# INLINE maxRight #-} -maxRight :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> (a -> Bool) -> m Int +maxRight :: + (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => + -- | The segment tree + SegTree (PrimState m) a -> + -- | \(l\) + Int -> + -- | \(p\): user prediate + (a -> Bool) -> + -- | \(r\): \(p\) holds for \([l, r)\) + m Int maxRight seg l0 f = maxRightM seg l0 (pure . f) --- | Moandic version of `maxRight`. +-- | Moandic variant of `maxRight`. -- -- ==== Constraints -- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side effect. @@ -400,7 +448,16 @@ -- -- @since 1.0.0.0 {-# INLINE maxRightM #-} -maxRightM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int +maxRightM :: + (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => + -- | The segment tree + SegTree (PrimState m) a -> + -- | \(l\) + Int -> + -- | \(p\): user prediate + (a -> m Bool) -> + -- | \(r\): \(p\) holds for \([l, r)\) + m Int maxRightM SegTree {..} l0 f = do b <- f mempty let !_ = ACIA.runtimeAssert b "AtCoder.SegTree.maxRightM: `f mempty` returned `False`" diff --git a/src/AtCoder/String.hs b/src/AtCoder/String.hs --- a/src/AtCoder/String.hs +++ b/src/AtCoder/String.hs @@ -1,9 +1,9 @@ -- | It contains string algorithms. -- --- Let @s@ be a string. We denote the substring of @s@ between \(a\)-th and \(b - 1\)-th character --- by @s[a..b)@. +-- Let \(s\) be a string. We denote the substring of \(s\) between \(a\)-th and \(b - 1\)-th +-- character by \(s[a..b)\). -- --- ==== Examples +-- ==== __Examples__ -- -- ===== Suffix Array and LCP Array -- @@ -32,9 +32,9 @@ -- * LCP array lcpArray, lcpArrayBS, - zAlgorithm, -- * Z algorithm + zAlgorithm, zAlgorithmBS, ) where @@ -55,9 +55,9 @@ -- | Calculates suffix array for a `Int` vector. -- --- Given a string @s@ of length \(n\), it returns the suffix array of @s@. Here, the suffix array --- @sa@ of @s@ is a permutation of \(0, \cdots, n-1\) such that @s[sa[i]..n) < s[sa[i+1]..n)@ holds --- for all \(i = 0,1, \cdots ,n-2\). +-- Given a string \(s\) of length \(n\), it returns the suffix array of \(s\). Here, the suffix array +-- \(\mathrm{sa}\) of \(s\) is a permutation of \(0, \cdots, n-1\) such that \(s[\mathrm{sa}[i]..n) < s[\mathrm{sa}[i+1]..n)\) +-- holds for all \(i = 0,1, \cdots ,n-2\). -- -- ==== Constraints -- - \(0 \leq n\) @@ -65,7 +65,7 @@ -- - \(0 \leq x \leq \mathrm{upper}\) for all elements \(x\) of \(s\). -- -- ==== Complexity --- - (3) \(O(n + \mathrm{upper})\)-time +-- - \(O(n + \mathrm{upper})\)-time -- -- @since 1.0.0.0 {-# INLINE suffixArray #-} @@ -81,7 +81,7 @@ -- - \(0 \leq n\) -- -- ==== Complexity --- - (1) \(O(n)\)-time +-- - \(O(n)\)-time -- -- @since 1.0.0.0 {-# INLINE suffixArrayBS #-} @@ -97,7 +97,8 @@ -- - \(0 \leq n\) -- -- ==== Complexity --- - (2) \(O(n \log n)\)-time, \(O(n)\)-space +-- - \(O(n \log n)\)-time +-- - \(O(n)\)-space -- -- @since 1.0.0.0 {-# INLINE suffixArrayOrd #-} @@ -124,12 +125,12 @@ (upper_,) <$> VU.unsafeFreeze vec in ACIS.saIs s2 upper --- | Given a string @s@ of length \(n\), it returns the LCP array of @s@. Here, the LCP array of --- @s@ is the array of length \(n-1\), such that the \(i\)-th element is the length of the LCP --- (Longest Common Prefix) of @s[sa[i]..n)@ and @s[sa[i+1]..n)@ +-- | Given a string \(s\) of length \(n\), it returns the LCP array of \(s\). Here, the LCP array of +-- \(s\) is the array of length \(n-1\), such that the \(i\)-th element is the length of the LCP +-- (Longest Common Prefix) of \(s[\mathrm{sa}[i]..n)\) and \(s[\mathrm{sa}[i+1]..n)\). -- -- ==== Constraints --- - The second argument is the suffix array of @s@. +-- - The second argument is the suffix array of \(s\). -- - \(1 \leq n\) -- -- ==== Complexity @@ -137,7 +138,14 @@ -- -- @since 1.0.0.0 {-# INLINE lcpArray #-} -lcpArray :: (HasCallStack, Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int -> VU.Vector Int +lcpArray :: + (HasCallStack, Ord a, VU.Unbox a) => + -- | A vector representing a string + VU.Vector a -> + -- | Suffix array + VU.Vector Int -> + -- | LCP array + VU.Vector Int lcpArray s sa = let n = VU.length s !_ = ACIA.runtimeAssert (n >= 1) "AtCoder.String.lcpArray: given empty input" @@ -167,10 +175,10 @@ rnk pure lcp --- | @ByteString@ verison of `lcpArray`. +-- | @ByteString@ variant of `lcpArray`. -- -- ==== Constraints --- - The second argument is the suffix array of @s@. +-- - The second argument is the suffix array of \(s\). -- - \(1 \leq n\) -- -- ==== Complexity @@ -178,17 +186,25 @@ -- -- @since 1.0.0.0 {-# INLINE lcpArrayBS #-} -lcpArrayBS :: (HasCallStack) => BS.ByteString -> VU.Vector Int -> VU.Vector Int +lcpArrayBS :: + (HasCallStack) => + -- | String + BS.ByteString -> + -- | Suffix array + VU.Vector Int -> + -- | LCP array + VU.Vector Int lcpArrayBS s sa = let n = BS.length s s2 = VU.map ord . VU.fromListN n $ BS.unpack s in lcpArray s2 sa -- | Given a `Ord` vector of length \(n\), it returns the array of length \(n\), such that the --- \(i\)-th element is the length of the LCP (Longest Common Prefix) of @s[0..n)@ and @s[i..n)@. +-- \(i\)-th element is the length of the LCP (Longest Common Prefix) of \(s[0..n)\) and \(s[i..n)\). -- -- ==== Constraints -- - \(n \leq n\) +-- -- ==== Complexity -- - \(O(n)\) -- @@ -227,10 +243,11 @@ n = VU.length s -- | Given a string of length \(n\), it returns the array of length \(n\), such that the \(i\)-th --- element is the length of the LCP (Longest Common Prefix) of @s[0..n)@ and @s[i..n)@. +-- element is the length of the LCP (Longest Common Prefix) of \(s[0..n)\) and \(s[i..n)\). -- -- ==== Constraints -- - \(n \leq n\) +-- -- ==== Complexity -- - \(O(n)\) -- diff --git a/src/AtCoder/TwoSat.hs b/src/AtCoder/TwoSat.hs --- a/src/AtCoder/TwoSat.hs +++ b/src/AtCoder/TwoSat.hs @@ -8,7 +8,7 @@ -- -- it decides whether there is a truth assignment that satisfies all clauses. -- --- ==== Example +-- ==== __Example__ -- >>> import AtCoder.TwoSat qualified as TS -- >>> import Data.Bit (Bit(..)) -- >>> ts <- TS.new 1 diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -4,8 +4,20 @@ import Test.Tasty.Ingredients.Rerun import Tests.Convolution qualified import Tests.Dsu qualified +import Tests.Extra.Bisect qualified +import Tests.Extra.HashMap qualified +import Tests.Extra.IntMap qualified +import Tests.Extra.IntSet qualified +import Tests.Extra.IntervalMap qualified import Tests.Extra.Math qualified import Tests.Extra.Monoid qualified +import Tests.Extra.MultiSet qualified +import Tests.Extra.Semigroup.Matrix qualified +import Tests.Extra.Semigroup.Permutation qualified +import Tests.Extra.WaveletMatrix qualified +import Tests.Extra.WaveletMatrix.BitVector qualified +import Tests.Extra.WaveletMatrix.Raw qualified +import Tests.Extra.WaveletMatrix2d qualified import Tests.FenwickTree qualified import Tests.Internal.Bit qualified import Tests.Internal.Buffer qualified @@ -31,8 +43,23 @@ . testGroup "toplevel" $ [ testGroup "Convolution" Tests.Convolution.tests, testGroup "Dsu" Tests.Dsu.tests, - testGroup "Extra.Math" Tests.Extra.Math.tests, - testGroup "Extra.Monoid" Tests.Extra.Monoid.tests, + testGroup + "Extra" + [ testGroup "Bisect" Tests.Extra.Bisect.tests, + testGroup "HashMap" Tests.Extra.HashMap.tests, + testGroup "IntervalMap" Tests.Extra.IntervalMap.tests, + testGroup "IntMap" Tests.Extra.IntMap.tests, + testGroup "IntSet" Tests.Extra.IntSet.tests, + testGroup "Math" Tests.Extra.Math.tests, + testGroup "Monoid" Tests.Extra.Monoid.tests, + testGroup "MultiSet" Tests.Extra.MultiSet.tests, + testGroup "Semigroup.Matrix" Tests.Extra.Semigroup.Matrix.tests, + testGroup "Semigroup.Permutation" Tests.Extra.Semigroup.Permutation.tests, + testGroup "WaveletMatrix" Tests.Extra.WaveletMatrix.tests, + testGroup "WaveletMatrix.BitVector" Tests.Extra.WaveletMatrix.BitVector.tests, + testGroup "WaveletMatrix.Raw" Tests.Extra.WaveletMatrix.Raw.tests, + testGroup "WaveletMatrix2d" Tests.Extra.WaveletMatrix2d.tests + ], testGroup "FenwickTree" Tests.FenwickTree.tests, testGroup "Internal" diff --git a/test/Tests/Extra/Bisect.hs b/test/Tests/Extra/Bisect.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/Bisect.hs @@ -0,0 +1,114 @@ +module Tests.Extra.Bisect where + +import AtCoder.Extra.Bisect +import Data.List qualified as L +import Data.Vector.Generic qualified as VG +import Data.Vector.Unboxed qualified as VU +import Test.Tasty +import Test.Tasty.QuickCheck as QC + +-- | Takes half-open interval [l, r). +naivePartition :: Int -> Int -> (Int -> Bool) -> VU.Vector Int -> (Maybe Int, Maybe Int) +naivePartition l r p xs + | l >= r = (Nothing, Nothing) + | otherwise = case (VU.null ls, VU.null rs) of + (True, True) -> error "unreachable" + (False, True) -> (Just l', Nothing) + (True, False) -> (Nothing, Just r') + _ -> (Just l', Just r') + where + xs' = VU.take (r - l) . VU.drop l $ xs + (!ls, !rs) = VU.partition p xs' + l' = l + VU.length ls - 1 + r' = l' + 1 + +naiveLowerBound :: VU.Vector Int -> Int -> Maybe Int +naiveLowerBound xs = naiveLowerBoundIn 0 (VU.length xs) xs + +naiveLowerBoundIn :: Int -> Int -> VU.Vector Int -> Int -> Maybe Int +naiveLowerBoundIn l r xs target = case naivePartition l r (< target) xs of + (!_, Just i) -> Just i + _ -> Nothing + +naiveUpperBound :: VU.Vector Int -> Int -> Maybe Int +naiveUpperBound xs = naiveUpperBoundIn 0 (VU.length xs) xs + +naiveUpperBoundIn :: Int -> Int -> VU.Vector Int -> Int -> Maybe Int +naiveUpperBoundIn l r xs target = case naivePartition l r (<= target) xs of + (!_, Just i) -> Just i + _ -> Nothing + +boundsQueryGen :: Gen (Int, Int, VU.Vector Int) +boundsQueryGen = do + n <- QC.chooseInt (1, 100) + p <- QC.chooseInt (-25, 25) + xs <- VU.fromList . L.sort <$> QC.vectorOf n (QC.chooseInt (-20, 20)) + pure (n, p, xs) + +bisectQueryGen :: Gen (Int, Int, VU.Vector Int, [(Int, Int)]) +bisectQueryGen = do + n <- QC.chooseInt (1, 100) + p <- QC.chooseInt (-25, 25) + xs <- VU.fromList . L.sort <$> QC.vectorOf n (QC.chooseInt (-20, 20)) + let lrs = [(l, r) | l <- [0 .. n], r <- [l .. n]] + pure (n, p, xs, lrs) + +prop_lowerBound :: TestTree +prop_lowerBound = QC.testProperty "lowerBound" $ do + (!_, !target, !xs) <- boundsQueryGen + pure $ naiveLowerBound xs target QC.=== lowerBound xs target + +prop_lowerBoundIn :: TestTree +prop_lowerBoundIn = QC.testProperty "lowerBoundIn" $ do + (!_, !target, !xs, !lrs) <- bisectQueryGen + pure . QC.conjoin $ + map + ( \(!l, !r) -> + naiveLowerBoundIn l r xs target == lowerBoundIn l r xs target + ) + lrs + +prop_upperBound :: TestTree +prop_upperBound = QC.testProperty "upperBound" $ do + (!_, !target, !xs) <- boundsQueryGen + pure $ naiveUpperBound xs target QC.=== upperBound xs target + +prop_upperBoundIn :: TestTree +prop_upperBoundIn = QC.testProperty "upperBoundIn" $ do + (!_, !target, !xs, !lrs) <- bisectQueryGen + pure . QC.conjoin $ + map + ( \(!l, !r) -> + naiveUpperBoundIn l r xs target == upperBoundIn l r xs target + ) + lrs + +prop_bisectL :: TestTree +prop_bisectL = QC.testProperty "bisectL" $ do + (!_, !boundary, !xs, !lrs) <- bisectQueryGen + pure . QC.conjoin $ + map + ( \(!l, !r) -> + fst (naivePartition l r (<= boundary) xs) == bisectL l r (\i -> xs VG.! i <= boundary) + ) + lrs + +prop_bisectR :: TestTree +prop_bisectR = QC.testProperty "bisectR" $ do + (!_, !boundary, !xs, !lrs) <- bisectQueryGen + pure . QC.conjoin $ + map + ( \(!l, !r) -> + snd (naivePartition l r (<= boundary) xs) == bisectR l r (\i -> xs VG.! i <= boundary) + ) + lrs + +tests :: [TestTree] +tests = + [ prop_lowerBound, + prop_upperBound, + prop_lowerBoundIn, + prop_upperBoundIn, + prop_bisectL, + prop_bisectR + ] diff --git a/test/Tests/Extra/Graph.hs b/test/Tests/Extra/Graph.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/Graph.hs @@ -0,0 +1,52 @@ +module Tests.Extra.Graph where + +import AtCoder.Extra.Graph qualified as Gr +import AtCoder.Internal.Buffer qualified as B +import Control.Monad (unless) +import Control.Monad.Fix (fix) +import Control.Monad.ST (runST) +import Data.List qualified as L +import Data.Vector.Generic qualified as VG +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import Test.Tasty +import Test.Tasty.QuickCheck as QC + +genDag :: Int -> QC.Gen (Gr.Csr ()) +genDag n = do + edges <- VU.fromList <$> QC.sublistOf [(u, v) | u <- [0 .. n - 1], v <- [u + 1 .. n - 1]] + verts <- VU.fromList <$> QC.shuffle [0 .. n - 1] + pure $ Gr.build n $ VU.map (\(!u, !v) -> (verts VG.! u, verts VG.! v, ())) edges + +dfs :: Int -> (Int -> VU.Vector Int) -> Int -> VU.Vector Int +dfs n gr u0 = runST $ do + buf <- B.new n + vis <- VUM.replicate n False + flip fix u0 $ \loop u -> do + VU.forM_ (gr u) $ \v -> do + b <- VUM.read vis v + unless b $ do + B.pushBack buf v + loop v + B.unsafeFreeze buf + +testTopSort :: Int -> Gr.Csr () -> VU.Vector Int -> Bool +testTopSort n gr vs = and + [ VU.notElem v (dfs n (gr `Gr.adj`) u) + | u <- (vs VG.!) <$> [0 .. n - 1], + v <- (vs VG.!) <$> [u + 1 .. n - 1] + ] + +-- | Tests lexicographically smallest topological ordering. +prop_topSort :: QC.Gen QC.Property +prop_topSort = do + n <- QC.chooseInt (1, 8) + dag <- genDag n + let vs = Gr.topSort n (dag `Gr.adj`) + let perms = map (VU.fromListN n) $ L.permutations [0 .. n - 1] + pure $ vs QC.=== head (filter (testTopSort n dag) perms) + +tests :: [TestTree] +tests = + [ QC.testProperty "topSort" prop_topSort + ] diff --git a/test/Tests/Extra/HashMap.hs b/test/Tests/Extra/HashMap.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/HashMap.hs @@ -0,0 +1,171 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.HashMap where + +import AtCoder.Extra.HashMap qualified as HM +import Control.Monad (foldM_) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (RealWorld) +import Data.HashMap.Strict qualified as HMR -- R: referencial implementation +import Data.Vector.Algorithms.Intro qualified as VAI +import Data.Vector.Unboxed qualified as VU +import System.IO.Unsafe (unsafePerformIO) +import Test.Hspec +import Test.QuickCheck.Monadic as QCM +import Test.Tasty +import Test.Tasty.Hspec +import Test.Tasty.QuickCheck as QC + +spec_invalid :: IO TestTree +spec_invalid = testSpec "capacity limit" $ do + it "throws error 1" $ do + hm <- HM.new @_ @Int 1 + HM.insert hm 0 0 + HM.insert hm 0 1 + HM.insert hm 1 2 `shouldThrow` anyException + + it "throws error 2" $ do + hm <- HM.new @_ @Int 2 + HM.insert hm 0 0 + HM.insert hm 1 1 + HM.insert hm 1 2 + HM.insert hm 2 2 `shouldThrow` anyException + +data Init = Init + { capacity :: {-# UNPACK #-} !Int, + ref0 :: !(HMR.HashMap Int Int), + hmM :: !(IO (HM.HashMap RealWorld Int)) + } + +instance Show Init where + show Init {..} = show (capacity, ref0) + +instance QC.Arbitrary Init where + arbitrary = do + capacity <- QC.chooseInt (1, 10) + pure $ Init capacity HMR.empty (HM.new capacity) + +data Query + = Size + | Member Int + | NotMember Int + | Lookup Int + | Insert Int Int + | InsertWithAdd Int Int + | Exchange Int Int + | ModifyAdd Int Int + | Clear + deriving (Show) + +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 + ] + where + keyGen = QC.chooseInt (-5, 5) + valGen = QC.chooseInt (-10, 10) + +-- | Arbitrary return type for the `Query` result. +data Result + = None + | B Bool + | I Int + | M (Maybe Int) + deriving (Show, Eq) + +-- | containers. (referencial implementation) +handleRef :: HMR.HashMap Int Int -> Query -> (HMR.HashMap Int Int, Result) +handleRef hm q = case q of + Size -> (hm, I $ HMR.size hm) + Member k -> (hm, B $ HMR.member k hm) + NotMember k -> (hm, B . not $ HMR.member k hm) + Lookup k -> (hm, M $ HMR.lookup k hm) + Insert k v -> (HMR.insert k v hm, None) + InsertWithAdd k v -> (HMR.insertWith (+) k v hm, None) + Exchange k v -> (HMR.insert k v hm, M $ HMR.lookup k hm) + ModifyAdd k v -> (HMR.adjust (+ v) k hm, None) + -- Delete k -> (HMR.delete k hm, HMR.lookup k hm) + Clear -> (HMR.empty, None) + +-- | ac-library-hs. +handleAcl :: (PrimMonad m) => HM.HashMap (PrimState m) Int -> Query -> m Result +handleAcl hm q = case q of + Size -> I <$> HM.size hm + Member k -> B <$> HM.member hm k + NotMember k -> B <$> HM.notMember hm k + Lookup k -> M <$> HM.lookup hm k + Insert k v -> do + HM.insert hm k v + pure None + InsertWithAdd k v -> do + HM.insertWith hm (+) k v + pure None + Exchange k v -> M <$> HM.exchange hm k v + ModifyAdd k v -> do + HM.modify hm (+ v) k + pure None + -- Delete k -> HM.delete hm k + Clear -> do + HM.clear hm + pure None + +-- | Ensures the capacity limit. +passQuery :: Int -> HMR.HashMap Int Int -> Query -> Bool +passQuery limit is (Insert k _) = HMR.member k is || HMR.size is < limit +passQuery limit is (InsertWithAdd k _) = HMR.member k is || HMR.size is < limit +passQuery limit is (Exchange k _) = HMR.member k is || HMR.size is < limit +passQuery _ _ _ = True + +prop_randomTest :: Init -> QC.Property +prop_randomTest Init {..} = QCM.monadicIO $ do + hm <- QCM.run hmM + q <- QCM.pick $ QC.chooseInt (1, 5 * capacity) + qs <- QCM.pick $ QC.vectorOf q (QC.arbitrary @Query) + foldM_ + ( \ref query -> do + if passQuery capacity ref query + then do + let (!ref', !expected) = handleRef ref query + res <- QCM.run $ handleAcl hm query + QCM.assertWith (expected == res) $ show (query, expected, res) + + -- check the map contents: + let assocsE = VU.modify (VAI.sortBy compare) . VU.fromList $ HMR.toList ref' + assocs <- QCM.run $ VU.modify (VAI.sortBy compare) <$> HM.unsafeAssocs hm + QCM.assertWith (assocsE == assocs) $ show ("- assocs", assocsE, assocs) + + let keysE = VU.modify (VAI.sortBy compare) . VU.fromList $ HMR.keys ref' + keys <- QCM.run $ VU.modify (VAI.sortBy compare) <$> HM.unsafeKeys hm + QCM.assertWith (keysE == keys) $ show ("- keys", keysE, keys) + + let elemsE = VU.modify (VAI.sortBy compare) . VU.fromList $ HMR.elems ref' + elems <- QCM.run $ VU.modify (VAI.sortBy compare) <$> HM.unsafeElems hm + QCM.assertWith (elemsE == elems) $ show ("- elems", elemsE, elems) + + let sizeE = HMR.size ref' + size <- QCM.run $ HM.size hm + QCM.assertWith (sizeE == size) $ show ("- size", sizeE, size) + + pure ref' + else pure ref + ) + ref0 + qs + +tests :: [TestTree] +tests = + [ unsafePerformIO spec_invalid, + QC.testProperty "random test" prop_randomTest + ] diff --git a/test/Tests/Extra/IntMap.hs b/test/Tests/Extra/IntMap.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/IntMap.hs @@ -0,0 +1,176 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.IntMap where + +import AtCoder.Extra.IntMap qualified as IM +import Control.Monad (foldM_) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (RealWorld) +import Data.Map.Strict qualified as IMR -- R: referencial implementation +import Data.Vector.Unboxed qualified as VU +import GHC.Stack (HasCallStack) +import Test.QuickCheck.Monadic as QCM +import Test.Tasty +import Test.Tasty.QuickCheck as QC + +data Init = Init + { n :: {-# UNPACK #-} !Int, + ref0 :: !(IMR.Map Int Int), + imM :: !(IO (IM.IntMap RealWorld Int)) + } + +instance Show Init where + show Init {..} = show (n, ref0) + +instance QC.Arbitrary Init where + arbitrary = do + n <- QC.chooseInt (1, 10) + pure $ Init n IMR.empty (IM.new n) + +data Query + = Size + | Member Int + | NotMember Int + | Null + | Lookup Int + | Insert Int Int + | InsertWithAdd Int Int + | Delete Int + | Delete_ Int + | LookupGE Int + | LookupGT Int + | LookupLE Int + | LookupLT Int + | LookupMin + | LookupMax + | DeleteMin + | DeleteMax + deriving (Show) + +-- | Arbitrary return type for the `Query` result. +data Result + = None + | B Bool + | I Int + | M (Maybe Int) + | KV (Maybe (Int, Int)) + | XS (VU.Vector Int) + | KVS (VU.Vector (Int, Int)) + deriving (Show, Eq) + +queryGen :: Int -> QC.Gen Query +queryGen n = do + QC.oneof + [ pure Size, + Member <$> lookupKeyGen, + NotMember <$> lookupKeyGen, + pure Null, + Lookup <$> lookupKeyGen, + -- insert is partial function + Insert <$> insertKeyGen <*> valGen, + InsertWithAdd <$> insertKeyGen <*> valGen, + Delete <$> lookupKeyGen, + Delete_ <$> lookupKeyGen, + LookupGE <$> lookupKeyGen, + LookupGT <$> lookupKeyGen, + LookupLE <$> lookupKeyGen, + LookupLT <$> lookupKeyGen, + pure LookupMin, + pure LookupMax + ] + where + -- for total functions + lookupKeyGen = QC.chooseInt (-1, n) + -- for partial functions + insertKeyGen = QC.chooseInt (0, n - 1) + valGen = QC.chooseInt (-10, 10) + +-- | containers. (referencial implementation) +handleRef :: IMR.Map Int Int -> Query -> (IMR.Map Int Int, Result) +handleRef im q = case q of + Size -> (im, I $ IMR.size im) + Member k -> (im, B $ IMR.member k im) + NotMember k -> (im, B . not $ IMR.member k im) + Null -> (im, B $ IMR.null im) + Lookup k -> (im, M $ IMR.lookup k im) + Insert k v -> (IMR.insert k v im, None) + InsertWithAdd k v -> (IMR.insertWith (+) k v im, None) + Delete k -> (IMR.delete k im, B $ IMR.member k im) + Delete_ k -> (IMR.delete k im, None) + LookupGE k -> (im, KV (IMR.lookupGE k im)) + LookupGT k -> (im, KV (IMR.lookupGT k im)) + LookupLE k -> (im, KV (IMR.lookupLE k im)) + LookupLT k -> (im, KV (IMR.lookupLT k im)) + LookupMin -> (im, KV (IMR.lookupMin im)) + LookupMax -> (im, KV (IMR.lookupMax im)) + DeleteMin -> wrapKV IMR.deleteFindMin + DeleteMax -> wrapKV IMR.deleteFindMax + where + wrapKV f + | IMR.null im = (im, KV Nothing) + | otherwise = let (!kv, !im') = f im in (im', KV (Just kv)) + +-- | ac-library-hs. +handleAcl :: (HasCallStack, PrimMonad m) => IM.IntMap (PrimState m) Int -> Query -> m Result +handleAcl im q = case q of + Size -> I <$> IM.size im + Member k -> B <$> IM.member im k + NotMember k -> B <$> IM.notMember im k + Null -> B <$> IM.null im + Lookup k -> M <$> IM.lookup im k + Insert k v -> do + IM.insert im k v + pure None + InsertWithAdd k v -> do + IM.insertWith im (+) k v + pure None + Delete k -> B <$> IM.delete im k + Delete_ k -> do + IM.delete_ im k + pure None + LookupGE k -> KV <$> IM.lookupGE im k + LookupGT k -> KV <$> IM.lookupGT im k + LookupLE k -> KV <$> IM.lookupLE im k + LookupLT k -> KV <$> IM.lookupLT im k + LookupMin -> KV <$> IM.lookupMin im + LookupMax -> KV <$> IM.lookupMax im + DeleteMin -> KV <$> IM.deleteMin im + DeleteMax -> KV <$> IM.deleteMax im + +prop_randomTest :: Init -> QC.Property +prop_randomTest Init {..} = QCM.monadicIO $ do + im <- QCM.run imM + q <- QCM.pick $ QC.chooseInt (1, 5 * n) + qs <- QCM.pick $ QC.vectorOf q (queryGen n) + foldM_ + ( \ref query -> do + let (!ref', !expected) = handleRef ref query + res <- QCM.run $ handleAcl im query + QCM.assertWith (expected == res) $ show (query, expected, res) + + -- check the map contents: + let assocsE = VU.fromList $ IMR.assocs ref' + assocs <- QCM.run $ IM.assocs im + QCM.assertWith (assocsE == assocs) $ show ("- assocs", assocsE, assocs) + + let keysE = VU.fromList $ IMR.keys ref' + keys <- QCM.run $ IM.keys im + QCM.assertWith (keysE == keys) $ show ("- keys", keysE, keys) + + let elemsE = VU.fromList $ IMR.elems ref' + elems <- QCM.run $ IM.elems im + QCM.assertWith (elemsE == elems) $ show ("- elems", elemsE, elems) + + let sizeE = IMR.size ref' + size <- QCM.run $ IM.size im + QCM.assertWith (sizeE == size) $ show ("- size", sizeE, size) + + pure ref' + ) + ref0 + qs + +tests :: [TestTree] +tests = + [ QC.testProperty "randomTest" prop_randomTest + ] diff --git a/test/Tests/Extra/IntSet.hs b/test/Tests/Extra/IntSet.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/IntSet.hs @@ -0,0 +1,150 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.IntSet where + +import AtCoder.Extra.IntSet qualified as IS +import Control.Monad (foldM_) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (RealWorld) +import Data.Set qualified as ISR -- R: referencial implementation +import Data.Vector.Unboxed qualified as VU +import GHC.Stack (HasCallStack) +import Test.QuickCheck.Monadic as QCM +import Test.Tasty +import Test.Tasty.QuickCheck as QC + +data Init = Init + { n :: {-# UNPACK #-} !Int, + ref0 :: !(ISR.Set Int), + isM :: !(IO (IS.IntSet RealWorld)) + } + +instance Show Init where + show Init {..} = show (n, ref0) + +instance QC.Arbitrary Init where + arbitrary = do + n <- QC.chooseInt (1, 10) + pure $ Init n ISR.empty (IS.new n) + +data Query + = Member Int + | NotMember Int + | Null + | Insert Int + | Delete Int + | Delete_ Int + | LookupGE Int + | LookupGT Int + | LookupLE Int + | LookupLT Int + | LookupMin + | LookupMax + | DeleteMin + | DeleteMax + deriving (Show) + +-- | Arbitrary return type for the `Query` result. +data Result + = None + | B Bool + | I Int + | M (Maybe Int) + deriving (Show, Eq) + +queryGen :: Int -> QC.Gen Query +queryGen n = do + QC.oneof + [ Member <$> lookupKeyGen, + NotMember <$> lookupKeyGen, + pure Null, + -- insert is partial function + Insert <$> insertKeyGen, + Delete <$> lookupKeyGen, + Delete_ <$> lookupKeyGen, + LookupGE <$> lookupKeyGen, + LookupGT <$> lookupKeyGen, + LookupLE <$> lookupKeyGen, + LookupLT <$> lookupKeyGen, + pure LookupMin, + pure LookupMax + ] + where + -- for total functions + lookupKeyGen = QC.chooseInt (-1, n) + -- for partial functions + insertKeyGen = QC.chooseInt (0, n - 1) + +-- | containers. (referencial implementation) +handleRef :: ISR.Set Int -> Query -> (ISR.Set Int, Result) +handleRef is q = case q of + Member k -> (is, B $ ISR.member k is) + NotMember k -> (is, B . not $ ISR.member k is) + Null -> (is, B $ ISR.null is) + Insert k -> (ISR.insert k is, None) + Delete k -> (ISR.delete k is, B $ ISR.member k is) + Delete_ k -> (ISR.delete k is, None) + LookupGE k -> (is, M (ISR.lookupGE k is)) + LookupGT k -> (is, M (ISR.lookupGT k is)) + LookupLE k -> (is, M (ISR.lookupLE k is)) + LookupLT k -> (is, M (ISR.lookupLT k is)) + LookupMin -> (is, M (ISR.lookupMin is)) + LookupMax -> (is, M (ISR.lookupMax is)) + DeleteMin -> wrapK ISR.deleteFindMin + DeleteMax -> wrapK ISR.deleteFindMax + where + wrapK f + | ISR.null is = (is, M Nothing) + | otherwise = let (!kv, !is') = f is in (is', M (Just kv)) + +-- | ac-library-hs. +handleAcl :: (HasCallStack, PrimMonad m) => IS.IntSet (PrimState m) -> Query -> m Result +handleAcl is q = case q of + Member k -> B <$> IS.member is k + NotMember k -> B <$> IS.notMember is k + Null -> B <$> IS.null is + Insert k -> do + IS.insert is k + pure None + Delete k -> B <$> IS.delete is k + Delete_ k -> do + IS.delete_ is k + pure None + LookupGE k -> M <$> IS.lookupGE is k + LookupGT k -> M <$> IS.lookupGT is k + LookupLE k -> M <$> IS.lookupLE is k + LookupLT k -> M <$> IS.lookupLT is k + LookupMin -> M <$> IS.lookupMin is + LookupMax -> M <$> IS.lookupMax is + DeleteMin -> M <$> IS.deleteMin is + DeleteMax -> M <$> IS.deleteMax is + +prop_randomTest :: Init -> QC.Property +prop_randomTest Init {..} = QCM.monadicIO $ do + is <- QCM.run isM + q <- QCM.pick $ QC.chooseInt (1, 5 * n) + qs <- QCM.pick $ QC.vectorOf q (queryGen n) + foldM_ + ( \ref query -> do + let (!ref', !expected) = handleRef ref query + res <- QCM.run $ handleAcl is query + QCM.assertWith (expected == res) $ show (query, expected, res) + + -- check the set contents: + let keysE = VU.fromList $ ISR.elems ref' + keys <- QCM.run $ IS.keys is + QCM.assertWith (keysE == keys) $ show ("- keys", keysE, keys) + + let sizeE = ISR.size ref' + size <- QCM.run $ IS.size is + QCM.assertWith (sizeE == size) $ show ("- size", sizeE, size) + + pure ref' + ) + ref0 + qs + +tests :: [TestTree] +tests = + [ QC.testProperty "randomTest" prop_randomTest + ] diff --git a/test/Tests/Extra/IntervalMap.hs b/test/Tests/Extra/IntervalMap.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/IntervalMap.hs @@ -0,0 +1,209 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.IntervalMap where + +import AtCoder.Extra.IntervalMap qualified as ITM +import AtCoder.Internal.Buffer qualified as ACIB +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (RealWorld, runST) +import Data.Foldable (for_) +import Data.List qualified as L +import Data.Map qualified as M +import Data.Primitive.MutVar +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import Test.QuickCheck.Monadic as QCM +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck as QC +import Tests.Util (intervalGen) + +-- | buildM should call onAdd and onDel correctly +prop_buildM :: QC.Gen QC.Property +prop_buildM = do + n <- QC.chooseInt (1, 10) + xs <- QC.vectorOf n (QC.chooseInt (-1, 1)) + + let groups = VU.group $ VU.fromList xs + let lens = L.scanl' (+) (0 :: Int) $ map VU.length groups + let expected = VU.fromList $ L.zip3 lens (tail lens) (L.map VU.head groups) + + let res = runST $ do + buf <- ACIB.new @_ @(Int, Int, Int) n + _ <- ITM.buildM (VU.fromList xs) $ \l r x -> do + ACIB.pushBack buf (l, r, x) + ACIB.freeze buf + + pure $ res QC.=== expected + +data Init = Init + { n :: {-# UNPACK #-} !Int, + refM :: !(IO (VUM.MVector RealWorld Int)), + itmM :: !(IO (ITM.IntervalMap RealWorld Int)) + } + +instance Show Init where + show Init {..} = show ("Init", n) + +instance QC.Arbitrary Init where + arbitrary = do + n <- QC.chooseInt (1, 20) + pure $ Init n (VUM.replicate n undef) (ITM.new n) + +data Query + = Contains Int + | Intersects (Int, Int) + | Lookup (Int, Int) + | -- | Read (Int, Int) + ReadMaybe (Int, Int) + | Insert (Int, Int) Int + | Delete (Int, Int) + | Overwrite (Int, Int) Int + | Freeze + | TestFreq + deriving (Show) + +-- | Arbitrary return type for the `Query` result. +data Result + = None + | B Bool + | I Int + | LRX (Maybe (Int, Int, Int)) + | M (Maybe Int) + | XS (VU.Vector (Int, (Int, Int))) + | -- | Counts (len * (len + 1) / 2) for each interval and sum them up. + Freq (M.Map Int Int) + deriving (Show, Eq) + +queryGen :: Int -> QC.Gen Query +queryGen n = do + QC.oneof + [ Contains <$> keyGen, + Intersects <$> intervalGen n, + Lookup <$> intervalGen n, + ReadMaybe <$> intervalGen n, + Insert <$> intervalGen n <*> valGen, + Delete <$> intervalGen n, + Overwrite <$> intervalGen n <*> valGen, + pure Freeze + -- TestFreq is manually given + -- pure TestFreq + ] + where + keyGen = QC.chooseInt (0, n - 1) + valGen = QC.chooseInt (-20, 20) + +undef :: Int +undef = minBound `div` 2 + +-- | Half-open intervals. +toIntervals :: (PrimMonad m) => VUM.MVector (PrimState m) Int -> m (VU.Vector (Int, Int, Int)) +toIntervals vec = do + vec' <- VU.freeze vec + let groups = VU.group vec' + let lens = L.scanl' (+) (0 :: Int) $ map VU.length groups + let intervals = zipWith (\xs l -> (l, l + VU.length xs, VU.head xs)) groups lens + pure . VU.fromList $ filter (\(!_, !_, !x) -> x /= undef) intervals + +-- | containers. (referencial implementation) +handleRef :: (PrimMonad m) => VUM.MVector (PrimState m) Int -> Query -> m Result +handleRef vec q = do + intervals <- toIntervals vec + case q of + Contains i -> do + pure . B $ VU.any (\(!l, !r, !_) -> l <= i && i < r) intervals + Intersects (!l, !r) + | l >= r -> pure $ B False + | otherwise -> pure . B $ VU.any (\(!l', !r', !_) -> l' <= l && r <= r') intervals + Lookup (!l, !r) + | l >= r -> pure $ LRX Nothing + | otherwise -> pure . LRX $ VU.find (\(!l', !r', !_) -> l' <= l && r <= r') intervals + ReadMaybe (!l, !r) + | l >= r -> pure $ M Nothing + | otherwise -> pure $ maybe (M Nothing) (M . Just . (\(!_, !_, !x) -> x)) $ VU.find (\(!l', !r', !_) -> l' <= l && r <= r') intervals + Insert (!l, !r) x -> do + for_ [l .. r - 1] $ \i -> do + VGM.write vec i x + pure None + Delete (!l, !r) -> do + for_ [l .. r - 1] $ \i -> do + VGM.write vec i undef + pure None + Overwrite (!l, !r) x + | l >= r -> pure None + | otherwise -> do + let interval = VU.find (\(!l', !r', !_) -> l' <= l && r <= r') intervals + case interval of + Just (!l', !r', !_) -> do + for_ [l' .. r' - 1] $ \i -> do + VGM.write vec i x + _ -> pure () + pure None + Freeze -> pure . XS $ VU.map (\(!l, !r, !x) -> (l, (r, x))) intervals + TestFreq -> pure . Freq . M.fromListWith (+) . VU.toList $ VU.map (\(!l, !r, !x) -> (x, (r - l) * ((r - l) + 1) `div` 2)) intervals + +-- | ac-library-hs. +handleAcl :: (HasCallStack, PrimMonad m) => MutVar (PrimState m) (M.Map Int Int) -> ITM.IntervalMap (PrimState m) Int -> Query -> m Result +handleAcl freq itm q = case q of + Contains i -> do + B <$> ITM.contains itm i + Intersects (!l, !r) -> do + B <$> ITM.intersects itm l r + Lookup (!l, !r) -> do + LRX <$> ITM.lookup itm l r + ReadMaybe (!l, !r) -> do + M <$> ITM.readMaybe itm l r + Insert (!l, !r) x -> do + ITM.insertM itm l r x onAdd onDel + pure None + Delete (!l, !r) -> do + ITM.deleteM itm l r onAdd onDel + pure None + Overwrite (!l, !r) x -> do + ITM.overwriteM itm l r x onAdd onDel + pure None + Freeze -> do + XS . VU.map (\(!l, (!r, !x)) -> (l, (r, x))) <$> ITM.freeze itm + TestFreq -> Freq <$> readMutVar freq + where + onAdd l r x = do + let len = r - l + let delta = len * (len + 1) `div` 2 + modifyMutVar freq $ M.insertWith (+) x delta + onDel l r x = modifyMutVar freq $ \m -> do + let len = r - l + let delta = len * (len + 1) `div` 2 + case M.lookup x m of + Just n + | n - delta == 0 -> M.delete x m + | otherwise -> M.insert x (n - delta) m + Nothing -> M.insert x delta m + +prop_randomTest :: Init -> QC.Property +prop_randomTest Init {..} = QCM.monadicIO $ do + ref <- QCM.run refM + im <- QCM.run itmM + freq <- QCM.run $ newMutVar M.empty + q <- QCM.pick $ QC.chooseInt (1, 10 * n) + qs <- QCM.pick $ QC.vectorOf q (queryGen n) + for_ qs $ \query -> do + expected <- QCM.run $ handleRef ref query + res <- QCM.run $ handleAcl freq im query + QCM.assertWith (expected == res) $ show (query, expected, res) + + -- always test Freq as an invariant + expectedI <- QCM.run $ handleRef ref Freeze + resI <- QCM.run $ handleAcl freq im Freeze + QCM.assertWith (expectedI == resI) $ show (Freeze, expectedI, resI) + + -- always test Freq as an invariant + expectedF <- QCM.run $ handleRef ref TestFreq + resF <- QCM.run $ handleAcl freq im TestFreq + QCM.assertWith (expectedF == resF) $ show (TestFreq, expectedF, resF) + +tests :: [TestTree] +tests = + [ QC.testProperty "buildM" prop_buildM, + QC.testProperty "randomTest" prop_randomTest + ] diff --git a/test/Tests/Extra/Monoid.hs b/test/Tests/Extra/Monoid.hs --- a/test/Tests/Extra/Monoid.hs +++ b/test/Tests/Extra/Monoid.hs @@ -1,7 +1,11 @@ module Tests.Extra.Monoid (tests) where import AtCoder.Extra.Monoid -import Data.Bit (Bit(..)) +import AtCoder.Extra.Monoid.Affine1 qualified as A +import AtCoder.Extra.Monoid.Mat2x2 qualified as M +import AtCoder.Extra.Monoid.RollingHash (RollingHash (..)) +import AtCoder.ModInt qualified as ModInt +import Data.Bit (Bit (..)) import Data.Proxy (Proxy (..)) import Data.Semigroup (Max (..), Min (..), Product (..), Sum (..), stimes) import Test.QuickCheck.Classes qualified as QCC @@ -71,10 +75,6 @@ arbitrary = RangeAdd <$> QC.arbitrary -- orphan instance -instance (QC.Arbitrary a) => QC.Arbitrary (RangeAddId a) where - arbitrary = RangeAddId <$> QC.arbitrary - --- orphan instance instance (QC.Arbitrary a, Monoid a) => QC.Arbitrary (RangeSet a) where arbitrary = do b <- (== 1) <$> QC.chooseInt (1, 30) @@ -83,12 +83,19 @@ else pure mempty -- orphan instance -instance (QC.Arbitrary a, Monoid a) => QC.Arbitrary (RangeSetId a) where +instance (QC.Arbitrary a) => QC.Arbitrary (Mat2x2 a) where + arbitrary = Mat2x2 <$> QC.arbitrary + +-- orphan instance +instance (QC.Arbitrary a) => QC.Arbitrary (V2 a) where + arbitrary = V2 <$> QC.arbitrary + +-- orphan instance +instance QC.Arbitrary (RollingHash b 998244353) where arbitrary = do - b <- (== 1) <$> QC.chooseInt (1, 30) - if b - then RangeSetId . (Bit True,) <$> QC.arbitrary - else pure mempty + hash <- QC.chooseInt (0, 998244353 - 1) + next <- QC.chooseInt (0, 998244353 - 1) + pure $ RollingHash hash next -- orphan instance instance QC.Arbitrary (Max Int) where @@ -98,11 +105,53 @@ instance QC.Arbitrary (Min Int) where arbitrary = Min <$> QC.arbitrary +-- orphan instance (TODO: move to common implementation) +instance QC.Arbitrary ModInt.ModInt998244353 where + arbitrary = ModInt.new <$> QC.arbitrary + +prop_affineZero :: Affine1 (Sum Int) -> QC.Property +prop_affineZero a = + QC.conjoin + [ A.zero <> a QC.=== A.zero, + a <> A.zero QC.=== (\(A.Affine1 (!_, !b)) -> A.Affine1 (0, b)) a + ] + +prop_affineIdent :: Affine1 (Sum Int) -> QC.Property +prop_affineIdent a = + QC.conjoin + [ A.ident <> a QC.=== a, + a <> A.ident QC.=== a + ] + +prop_mat2x2Zero :: Mat2x2 Int -> QC.Property +prop_mat2x2Zero a = + QC.conjoin + [ M.zero <> a QC.=== M.zero, + a <> M.zero QC.=== M.zero + ] + +prop_mat2x2Ident :: Mat2x2 Int -> QC.Property +prop_mat2x2Ident a = + QC.conjoin + [ M.ident <> a QC.=== a, + a <> M.ident QC.=== a + ] + +prop_mat2x2Inv :: Mat2x2 ModInt.ModInt998244353 -> QC.Property +prop_mat2x2Inv a = + (M.det a /= 0 QC.==>) $ + QC.conjoin + [ M.inv a <> a QC.=== M.ident, + a <> M.inv a QC.=== M.ident + ] + tests :: [TestTree] tests = [ testGroup "Affine1" - [ laws @(Affine1 (Sum Int)) + [ QC.testProperty "zero" prop_affineZero, + QC.testProperty "ident" prop_affineIdent, + laws @(Affine1 (Sum Int)) [ QCC.semigroupLaws, QCC.monoidLaws, QCC.semigroupMonoidLaws @@ -113,26 +162,28 @@ ], testGroup "RangeAdd" - [ laws @(RangeAdd Int) + [ laws @(RangeAdd (Sum Int)) [ QCC.semigroupLaws, QCC.monoidLaws, QCC.semigroupMonoidLaws ], - laws @(RangeAdd Int, Sum Int) + laws @(RangeAdd (Sum Int), Sum Int) [ segActLaw - ] - ], - testGroup - "RangeAddId" - [ laws @(RangeAddId Int) + ], + laws @(RangeAdd (Max Int)) [ QCC.semigroupLaws, QCC.monoidLaws, QCC.semigroupMonoidLaws ], - laws @(RangeAddId Int, Max Int) + laws @(RangeAdd (Max Int), Max Int) [ segActLaw ], - laws @(RangeAddId Int, Min Int) + laws @(RangeAdd (Min Int)) + [ QCC.semigroupLaws, + QCC.monoidLaws, + QCC.semigroupMonoidLaws + ], + laws @(RangeAdd (Min Int), Min Int) [ segActLaw ] ], @@ -143,65 +194,62 @@ QCC.monoidLaws, QCC.semigroupMonoidLaws ], + laws @(RangeSet (Sum Int), Sum Int) + [ segActLaw + ], laws @(RangeSet (Product Int)) [ QCC.semigroupLaws, QCC.monoidLaws, QCC.semigroupMonoidLaws ], + laws @(RangeSet (Product Int), Product Int) + [ segActLaw + ], laws @(RangeSet (Max Int)) [ QCC.semigroupLaws, QCC.monoidLaws, QCC.semigroupMonoidLaws ], + laws @(RangeSet (Max Int), Max Int) + [ segActLaw + ], laws @(RangeSet (Min Int)) [ QCC.semigroupLaws, QCC.monoidLaws, QCC.semigroupMonoidLaws ], - laws @(RangeSet (Sum Int), Sum Int) - [ segActLaw - ], - laws @(RangeSet (Product Int), Product Int) - [ segActLaw - ], - laws @(RangeSet (Max Int), Max Int) - [ segActLaw - ], laws @(RangeSet (Min Int), Min Int) [ segActLaw ] ], testGroup - "RangeSetId" - [ laws @(RangeSetId (Sum Int)) + "Mat2x2" + [ QC.testProperty "zero" prop_mat2x2Zero, + QC.testProperty "ident" prop_mat2x2Ident, + QC.testProperty "inv" prop_mat2x2Inv, + laws @(Mat2x2 Int) [ QCC.semigroupLaws, - QCC.idempotentSemigroupLaws, QCC.monoidLaws, QCC.semigroupMonoidLaws ], - laws @(RangeSetId (Product Int)) + laws @(Mat2x2 Int, V2 Int) + [ segActLaw + ] + ], + testGroup + "V2" + [ laws @(V2 Int) [ QCC.semigroupLaws, - QCC.idempotentSemigroupLaws, QCC.monoidLaws, QCC.semigroupMonoidLaws - ], - laws @(RangeSetId (Max Int)) + ] + ], + testGroup + "RollingHash" + [ laws @(RollingHash 100 998244353) [ QCC.semigroupLaws, - QCC.idempotentSemigroupLaws, QCC.monoidLaws, QCC.semigroupMonoidLaws - ], - laws @(RangeSetId (Min Int)) - [ QCC.semigroupLaws, - QCC.idempotentSemigroupLaws, - QCC.monoidLaws, - QCC.semigroupMonoidLaws - ], - laws @(RangeSetId (Max Int), Max Int) - [ segActLaw - ], - laws @(RangeSetId (Min Int), Min Int) - [ segActLaw ] ] ] diff --git a/test/Tests/Extra/MultiSet.hs b/test/Tests/Extra/MultiSet.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/MultiSet.hs @@ -0,0 +1,218 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.MultiSet (tests) where + +import AtCoder.Extra.MultiSet qualified as MS +import Control.Monad (foldM_, when) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (RealWorld) +import Data.IntMap qualified as IM +import Data.IntSet qualified as IS +import Data.Vector.Algorithms.Intro qualified as VAI +import Data.Vector.Unboxed qualified as VU +import System.IO.Unsafe (unsafePerformIO) +import Test.Hspec +import Test.QuickCheck.Monadic as QCM +import Test.Tasty +import Test.Tasty.Hspec +import Test.Tasty.QuickCheck as QC + +spec_invalid :: IO TestTree +spec_invalid = testSpec "capacity limit" $ do + it "throws error 1" $ do + ms <- MS.new @_ 1 + MS.inc ms 0 + MS.inc ms 0 + MS.inc ms 1 `shouldThrow` anyException + + it "throws error 2" $ do + ms <- MS.new @_ 2 + MS.inc ms 0 + MS.inc ms 1 + MS.inc ms 2 `shouldThrow` anyException + + it "throws error 2" $ do + ms <- MS.new @_ 2 + MS.inc ms 0 + MS.inc ms 1 + MS.delete ms 1 + MS.inc ms 2 `shouldThrow` anyException + +data Init = Init + { n :: {-# UNPACK #-} !Int, + ref0 :: !(IM.IntMap Int), + msM :: !(IO (MS.MultiSet RealWorld)) + } + +instance Show Init where + show Init {..} = show (n, ref0) + +instance QC.Arbitrary Init where + arbitrary = do + n <- QC.chooseInt (1, 10) + pure $ Init n IM.empty (MS.new n) + +data Query + = Member Int + | NotMember Int + | Lookup Int + | Inc Int + | -- Safe* are performed to not create negative count + SafeDec Int + | SafeAdd Int Int + | SafeSub Int Int + | Insert Int Int + | Delete Int + deriving (Show) + +-- | Arbitrary return type for the `Query` result. +data Result + = None + | B Bool + | I Int + | M (Maybe Int) + deriving (Show, Eq) + +queryGen :: Int -> QC.Gen Query +queryGen n = do + QC.oneof + [ Member <$> keyGen, + NotMember <$> keyGen, + Lookup <$> keyGen, + Inc <$> keyGen, + SafeDec <$> keyGen, + SafeAdd <$> keyGen <*> deltaGen, + SafeSub <$> keyGen <*> deltaGen, + Insert <$> keyGen <*> insertValGen, + Delete <$> keyGen + ] + where + keyGen = QC.chooseInt (-n, n) + deltaGen = QC.chooseInt (-n, n) + insertValGen = QC.chooseInt (1, n) + +-- | containers. (referencial implementation) +handleRef :: IM.IntMap Int -> Query -> (IM.IntMap Int, Result) +handleRef im q = case q of + Member key -> (im, B $ IM.member key im) + NotMember key -> (im, B $ IM.notMember key im) + Lookup key -> (im, M $ IM.lookup key im) + Inc key -> (IM.insertWith (+) key 1 im, None) + SafeDec key -> case IM.lookup key im of + Just 1 -> (IM.delete key im, None) + Just n -> (IM.insert key (n - 1) im, None) + Nothing -> (im, None) + SafeAdd key dx -> case IM.lookup key im of + Just n | n + dx == 0 -> (IM.delete key im, None) + Just n | n + dx > 0 -> (IM.insert key (n + dx) im, None) + Nothing | dx > 0 -> (IM.insert key dx im, None) + _ -> (im, None) + SafeSub key dx -> case IM.lookup key im of + Just n | n - dx == 0 -> (IM.delete key im, None) + Just n | n - dx > 0 -> (IM.insert key (n - dx) im, None) + Nothing | -dx > 0 -> (IM.insert key (-dx) im, None) + _ -> (im, None) + Insert key val -> (IM.insert key val im, None) + Delete key -> (IM.delete key im, None) + +-- | ac-library-hs. +handleAcl :: (HasCallStack, PrimMonad m) => MS.MultiSet (PrimState m) -> Query -> m Result +handleAcl ms q = case q of + Member key -> B <$> MS.member ms key + NotMember key -> B <$> MS.notMember ms key + Lookup key -> M <$> MS.lookup ms key + Inc key -> do + MS.inc ms key + pure None + SafeDec key -> do + b <- MS.member ms key + when b $ do + MS.dec ms key + pure None + SafeAdd key dx -> do + MS.lookup ms key >>= \case + Just n | n + dx >= 0 -> do + MS.add ms key dx + Nothing | dx >= 0 -> do + MS.add ms key dx + _ -> pure () + pure None + SafeSub key dx -> do + MS.lookup ms key >>= \case + Just n | n - dx >= 0 -> do + MS.sub ms key dx + Nothing | -dx >= 0 -> do + MS.sub ms key dx + _ -> pure () + pure None + Insert key val -> do + MS.insert ms key val + pure None + Delete key -> do + MS.delete ms key + pure None + +-- | Ensures the capacity limit. +passQuery :: Int -> IS.IntSet -> Query -> Bool +passQuery limit is (Inc k) = IS.member k is || IS.size is < limit +passQuery limit is (SafeAdd k dx) = IS.member k is || IS.size is < limit || dx <= 0 +passQuery limit is (SafeSub k dx) = IS.member k is || IS.size is < limit || dx >= 0 +passQuery limit is (Insert k _) = IS.member k is || IS.size is < limit +passQuery _ _ _ = True + +-- | Records used keys. +recordKey :: IS.IntSet -> Query -> IS.IntSet +recordKey is (Inc k) = IS.insert k is +recordKey is (SafeAdd k dx) | dx > 0 = IS.insert k is +recordKey is (SafeSub k dx) | dx < 0 = IS.insert k is +recordKey is (Insert k _) = IS.insert k is +recordKey is _ = is + +-- TODO: record used key + +prop_randomTest :: Init -> QC.Property +prop_randomTest Init {..} = QCM.monadicIO $ do + ms <- QCM.run msM + q <- QCM.pick $ QC.chooseInt (1, 5 * n) + qs <- QCM.pick $ QC.vectorOf q (queryGen n) + + foldM_ + ( \(!ref, !keys) query -> do + if passQuery n keys query + then do + -- run the query + let (!ref', !expected) = handleRef ref query + res <- QCM.run $ handleAcl ms query + QCM.assertWith (expected == res) $ show (query, expected, res) + + -- check the map contents: + let assocsE = VU.modify (VAI.sortBy compare) (VU.fromList (IM.assocs ref')) + assocs <- QCM.run $ VU.modify (VAI.sortBy compare) <$> MS.unsafeAssocs ms + QCM.assertWith (assocsE == assocs) $ show ("- assocs", assocsE, assocs) + + let keysE = VU.modify (VAI.sortBy compare) (VU.fromList (IM.keys ref')) + keys_ <- QCM.run $ VU.modify (VAI.sortBy compare) <$> MS.unsafeKeys ms + QCM.assertWith (keysE == keys_) $ show ("- keys", keysE, keys_) + + let elemsE = VU.modify (VAI.sortBy compare) (VU.fromList (IM.elems ref')) + elems <- QCM.run $ VU.modify (VAI.sortBy compare) <$> MS.unsafeElems ms + QCM.assertWith (elemsE == elems) $ show ("- elems", elemsE, elems) + + let sizeE = IM.size ref' + size <- QCM.run $ MS.size ms + QCM.assertWith (sizeE == size) $ show ("- size", sizeE, size) + + pure (ref', recordKey keys query) + else + pure (ref, keys) + ) + (ref0, IS.empty) + qs + +-- TODO: test invariant + +tests :: [TestTree] +tests = + [ unsafePerformIO spec_invalid, + QC.testProperty "random test" prop_randomTest + ] diff --git a/test/Tests/Extra/Semigroup/Matrix.hs b/test/Tests/Extra/Semigroup/Matrix.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/Semigroup/Matrix.hs @@ -0,0 +1,43 @@ +module Tests.Extra.Semigroup.Matrix (tests) where + +import AtCoder.Extra.Semigroup.Matrix qualified as Mat +import AtCoder.ModInt qualified as M +import Data.Vector.Unboxed qualified as VU +import GHC.TypeNats (KnownNat) +import Test.QuickCheck.Classes qualified as QCC +import Test.Tasty +import Test.Tasty.QuickCheck qualified as QC +import Tests.Util (laws) + +-- TODO: (const True) should be removed + +-- orphan instance +instance (QC.Arbitrary a, VU.Unbox a) => QC.Arbitrary (Mat.Matrix a) where + -- for simplicity, make a 33x33 matrix + arbitrary = do + let n = 33 + vec <- VU.fromList <$> QC.vectorOf (n * n) (QC.arbitrary @a) + pure $ Mat.Matrix n n vec + +-- orphan instance +instance (KnownNat p) => QC.Arbitrary (M.ModInt p) where + arbitrary = M.new <$> QC.arbitrary + +prop_mulToCol :: QC.Gen QC.Property +prop_mulToCol = do + h <- QC.chooseInt (1, 16) + w <- QC.chooseInt (1, 16) + vec <- VU.fromList <$> QC.vectorOf (h * w) (QC.arbitrary @Int) + let mat = Mat.new h w vec + col <- VU.fromList <$> QC.vectorOf w (QC.arbitrary @Int) + let lhs = Mat.mulToCol mat col + let rhs = Mat.vecM $ Mat.mul mat (Mat.new w 1 col) + pure $ lhs QC.=== rhs + +tests :: [TestTree] +tests = + [ QC.testProperty "mulToCol" prop_mulToCol, + laws @(Mat.Matrix (M.ModInt 998244353)) + [ QCC.semigroupLaws + ] + ] diff --git a/test/Tests/Extra/Semigroup/Permutation.hs b/test/Tests/Extra/Semigroup/Permutation.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/Semigroup/Permutation.hs @@ -0,0 +1,71 @@ +module Tests.Extra.Semigroup.Permutation (tests) where + +import AtCoder.Extra.Semigroup.Permutation qualified as P +import Control.Exception (evaluate) +import Data.Vector.Unboxed qualified as VU +import System.IO.Unsafe (unsafePerformIO) +import Test.Hspec +import Test.QuickCheck.Classes qualified as QCC +import Test.Tasty +import Test.Tasty.Hspec +import Test.Tasty.QuickCheck as QC +import Tests.Util (laws) + +spec_invalid :: IO TestTree +spec_invalid = testSpec "boundary check" $ do + let !_ = P.new $ VU.fromList [0] + let !_ = P.new $ VU.fromList [-1] -- -1 is allowed + it "throws error 1" $ do + evaluate (P.new (VU.fromList [1])) `shouldThrow` anyException + it "throws error 2" $ do + evaluate (P.new (VU.fromList [-2])) `shouldThrow` anyException + +prop_ident :: P.Permutation -> QC.Property +prop_ident p = + QC.conjoin + [ p <> ident QC.=== p, + ident <> p QC.=== p + ] + where + ident = P.ident (P.length p) + +prop_zero :: P.Permutation -> QC.Property +prop_zero p = + QC.conjoin + [ p <> zero QC.=== zero, + zero <> p QC.=== zero + ] + where + zero = P.zero (P.length p) + +prop_identAct :: QC.Positive Int -> QC.Gen QC.Property +prop_identAct (QC.Positive len) = do + let p = P.ident len + i <- QC.chooseInt (0, len - 1) + pure $ P.act p i QC.=== i + +prop_zeroAct :: QC.Positive Int -> QC.Gen QC.Property +prop_zeroAct (QC.Positive len) = do + let p = P.zero len + i <- QC.chooseInt (0, len - 1) + pure $ P.act p i QC.=== i + +-- orphan instance +instance QC.Arbitrary P.Permutation where + arbitrary = do + let n = 33 + vec <- VU.fromList <$> QC.vectorOf n (QC.chooseInt (-1, n - 1)) + pure $ P.new vec + +tests :: [TestTree] +tests = + [ unsafePerformIO spec_invalid, + QC.testProperty "ident" prop_ident, + QC.testProperty "zero" prop_zero, + QC.testProperty "identAct" prop_identAct, + QC.testProperty "zeroAct" prop_zeroAct, + laws + @P.Permutation + [ QCC.semigroupLaws + ] + ] diff --git a/test/Tests/Extra/WaveletMatrix.hs b/test/Tests/Extra/WaveletMatrix.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/WaveletMatrix.hs @@ -0,0 +1,237 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.WaveletMatrix (tests) where + +import AtCoder.Extra.WaveletMatrix qualified as WM +import Data.IntMap qualified as IM +import Data.Ord (comparing) +import Data.Vector.Algorithms.Intro qualified as VAI +import Data.Vector.Unboxed qualified as VU +import Test.Hspec +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck as QC +import Tests.Util (intervalGen) + +data Init = Init + { capacity :: {-# UNPACK #-} !Int, + xs :: !(VU.Vector Int), + wm :: !WM.WaveletMatrix + } + +instance Show Init where + show Init {..} = show ("Init", capacity, xs) + +instance QC.Arbitrary Init where + arbitrary = do + QC.NonNegative n <- QC.arbitrary + xs <- VU.fromListN n <$> QC.vectorOf n (QC.arbitrary @Int) + pure $ Init n xs (WM.build xs) + +data Query + = Access !Int + | KthSmallestIn !(Int, Int) !Int + | IKthSmallestIn !(Int, Int) !Int + | KthLargestIn !(Int, Int) !Int + | IKthLargestIn !(Int, Int) !Int + | RankBetween !(Int, Int) !(Int, Int) + | Rank !(Int, Int) !Int + | Select !Int + | SelectIn !(Int, Int) !Int + | SelectKth !Int !Int + | SelectKthIn !(Int, Int) !Int !Int + | LookupLE !(Int, Int) !Int + | LookupLT !(Int, Int) !Int + | LookupGE !(Int, Int) !Int + | LookupGT !(Int, Int) !Int + | AssocsIn !(Int, Int) + | DescAssocsIn !(Int, Int) + deriving (Show) + +genQuery :: Int -> QC.Gen Query +genQuery n = do + QC.oneof + [ Access <$> QC.chooseInt (-1, n), + KthSmallestIn <$> lr <*> exc, + IKthSmallestIn <$> lr <*> exc, + KthLargestIn <$> lr <*> exc, + IKthLargestIn <$> lr <*> exc, + RankBetween <$> lr <*> lr, + Rank <$> lr <*> val, + Select <$> val, + SelectIn <$> lr <*> val, + SelectKth <$> exc <*> val, + SelectKthIn <$> lr <*> exc <*> val, + LookupLE <$> lr <*> QC.chooseInt (-1, n), + LookupLT <$> lr <*> QC.chooseInt (-1, n), + LookupGE <$> lr <*> QC.chooseInt (-1, n), + LookupGT <$> lr <*> QC.chooseInt (-1, n), + AssocsIn <$> lr, + DescAssocsIn <$> lr + ] + where + exc = QC.chooseInt (0, n) + lr = intervalGen n + val = QC.chooseInt (-n, n) + +-- | Arbitrary return type for the `Query` result. +data Result + = I {-# UNPACK #-} !Int + | M !(Maybe Int) + | M2 !(Maybe (Int, Int)) + | Assocs [(Int, Int)] + deriving (Show, Eq) + +-- | containers. (referencial implementation) +handleRef :: VU.Vector Int -> Query -> Result +handleRef xs q = case q of + Access i -> M $ xs VU.!? i + KthSmallestIn (!l, !r) k -> M $ snd <$> ikthSmallestIn l r k + IKthSmallestIn (!l, !r) k -> M2 $ ikthSmallestIn l r k + KthLargestIn (!l, !r) k -> M $ snd <$> ikthLargestIn l r k + IKthLargestIn (!l, !r) k -> M2 $ ikthLargestIn l r k + RankBetween (!l, !r) (!xl, !xr) -> rankBetween l r xl xr + Rank (!l, !r) x -> rankBetween l r x (x + 1) + Select x -> M $ selectKthIn 0 n 0 x + SelectIn (!l, !r) x -> M $ selectKthIn l r 0 x + SelectKth k x -> M $ selectKthIn 0 n k x + SelectKthIn (!l, !r) k x -> M $ selectKthIn l r k x + LookupLE (!l, !r) x -> max_ . VU.filter (<= x) . VU.take (r - l) $ VU.drop l xs + LookupLT (!l, !r) x -> max_ . VU.filter (< x) . VU.take (r - l) $ VU.drop l xs + LookupGE (!l, !r) x -> min_ . VU.filter (>= x) . VU.take (r - l) $ VU.drop l xs + LookupGT (!l, !r) x -> min_ . VU.filter (> x) . VU.take (r - l) $ VU.drop l xs + AssocsIn (!l, !r) -> Assocs . IM.assocs . IM.fromListWith (+) . VU.toList . VU.map (,1) . VU.take (r - l) $ VU.drop l xs + DescAssocsIn (!l, !r) -> Assocs . reverse . IM.assocs . IM.fromListWith (+) . VU.toList . VU.map (,1) . VU.take (r - l) $ VU.drop l xs + where + n = VU.length xs + max_ ys + | VU.null ys = M Nothing + | otherwise = M $ Just $ VU.maximum ys + min_ ys + | VU.null ys = M Nothing + | otherwise = M $ Just $ VU.minimum ys + ikthSmallestIn l r k = + (VU.!? k) + . VU.modify (VAI.sortBy (comparing (\(!i, !x) -> (x, i)))) + . VU.take (r - l) + . VU.drop l + . VU.indexed + $ xs + ikthLargestIn l r k = ikthSmallestIn l r ((r - l) - (k + 1)) + rankBetween l r xl xr = + I + . VU.length + . VU.filter (\x -> xl <= x && x < xr) + . VU.take (r - l) + . VU.drop l + $ xs + selectKthIn l r k x = + (fst <$>) + . (VU.!? k) + . VU.filter ((== x) . snd) + . VU.take (r - l) + . VU.drop l + . VU.indexed + $ xs + +handleAcl :: WM.WaveletMatrix -> Query -> Result +handleAcl wm q = case q of + Access i -> M $ WM.access wm i + KthSmallestIn (!l, !r) k -> M $ WM.kthSmallestIn wm l r k + IKthSmallestIn (!l, !r) k -> M2 $ WM.ikthSmallestIn wm l r k + KthLargestIn (!l, !r) k -> M $ WM.kthLargestIn wm l r k + IKthLargestIn (!l, !r) k -> M2 $ WM.ikthLargestIn wm l r k + RankBetween (!l, !r) (!xl, !xr) -> I $ WM.rankBetween wm l r xl xr + Rank (!l, !r) x -> I $ WM.rankBetween wm l r x (x + 1) + Select x -> M $ WM.select wm x + SelectIn (!l, !r) x -> M $ WM.selectIn wm l r x + SelectKth k x -> M $ WM.selectKth wm k x + SelectKthIn (!l, !r) k x -> M $ WM.selectKthIn wm l r k x + LookupLE (!l, !r) x -> M $ WM.lookupLE wm l r x + LookupLT (!l, !r) x -> M $ WM.lookupLT wm l r x + LookupGE (!l, !r) x -> M $ WM.lookupGE wm l r x + LookupGT (!l, !r) x -> M $ WM.lookupGT wm l r x + AssocsIn (!l, !r) -> Assocs $ WM.assocsIn wm l r + DescAssocsIn (!l, !r) -> Assocs $ WM.descAssocsIn wm l r + +prop_randomTest :: Init -> QC.Gen QC.Property +prop_randomTest Init {..} = do + qs <- QC.vectorOf capacity (genQuery capacity) + pure . QC.conjoin $ + map + ( \q -> + QC.counterexample (show q) $ + handleRef xs q QC.=== handleAcl wm q + ) + qs + +unit_boundary :: TestTree +unit_boundary = testCase "boundary" $ do + let n = 5 + let wm = WM.build $ VU.fromList [0, 10, 20, 10, 0] + + let try :: (HasCallStack, Eq a, Show a) => (WM.WaveletMatrix -> Int -> Int -> Int -> Maybe a) -> Int -> IO () + try f x = do + (@?= Nothing) $ f wm (-1) 0 x + (@?= Nothing) $ f wm 20 (20 + 1) x + + let k = 0 + try WM.kthSmallestIn k + try WM.ikthSmallestIn k + try WM.kthLargestIn k + try WM.ikthLargestIn k + +-- try WM.unsafeKthSmallestIn +-- try WM.unsafeIKthSmallestIn +-- try WM.unsafeKthLargestIn +-- try WM.unsafeIKthLargestIn + + let tryRank :: (HasCallStack) => (WM.WaveletMatrix -> Int -> Int -> Int -> Int) -> Int -> IO () + tryRank f x = do + -- out of range + (@?= 0) $ f wm (-1) 0 x + (@?= 0) $ f wm n (n + 1) x + -- reverse + (@?= 0) $ f wm 1 0 x + (@?= 0) $ f wm n (n - 1) x + (@?= 0) $ f wm n 0 x + -- out of range and reverse + (@?= 0) $ f wm 0 (-1) x + (@?= 0) $ f wm (n + 1) n x + + let tryRankBetween :: (HasCallStack) => (WM.WaveletMatrix -> Int -> Int -> Int -> Int -> Int) -> Int -> Int -> IO () + tryRankBetween f xl xr = do + -- out of range + (@?= 0) $ f wm (-1) 0 xl xr + (@?= 0) $ f wm n (n + 1) xl xr + -- reverse + (@?= 0) $ f wm 1 0 xl xr + (@?= 0) $ f wm n (n - 1) xl xr + (@?= 0) $ f wm n 0 xl xr + -- out of range and reverse + (@?= 0) $ f wm 0 (-1) xl xr + (@?= 0) $ f wm (n + 1) n xl xr + + tryRankBetween WM.rankBetween (-1) 1 + tryRank WM.rank 0 + tryRank WM.rank 1 + tryRank WM.rank (-1) + +-- TODO: test + +-- (@?= Nothing) $ WM.select wm +-- (@?= Nothing) $ WM.selectKth wm +-- (@?= Nothing) $ WM.selectKthIn wm + +-- (@?= Nothing) $ WM.lookupLE wm +-- (@?= Nothing) $ WM.lookupLT wm +-- (@?= Nothing) $ WM.lookupGE wm +-- (@?= Nothing) $ WM.lookupGT wm +-- (@?= Nothing) $ WM.assocsIn wm +-- (@?= Nothing) $ WM.descAssocsIn wm + +tests :: [TestTree] +tests = + [ unit_boundary, + QC.testProperty "random test" prop_randomTest + ] diff --git a/test/Tests/Extra/WaveletMatrix/BitVector.hs b/test/Tests/Extra/WaveletMatrix/BitVector.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/WaveletMatrix/BitVector.hs @@ -0,0 +1,84 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.WaveletMatrix.BitVector where + +import AtCoder.Extra.WaveletMatrix.BitVector qualified as BV +import Data.Bit (Bit (..)) +import Data.Vector.Unboxed qualified as VU +import Test.Tasty +import Test.Tasty.QuickCheck as QC +import Tests.Util (intervalGen) + +data Init = Init + { capacity :: {-# UNPACK #-} !Int, + ref :: !(VU.Vector Bool), + bv :: !BV.BitVector + } + deriving (Eq, Show) + +instance QC.Arbitrary Init where + arbitrary = do + QC.Positive capacity <- QC.arbitrary + bs <- QC.vectorOf capacity (QC.arbitrary @Bool) + pure $ Init capacity (VU.fromList bs) (BV.build (VU.fromList (map Bit bs))) + +data Query + = Rank0 {-# UNPACK #-} !Int + | Rank1 {-# UNPACK #-} !Int + | Select0 {-# UNPACK #-} !Int + | Select1 {-# UNPACK #-} !Int + | SelectKth0 {-# UNPACK #-} !(Int, Int) !Int + | SelectKth1 {-# UNPACK #-} !(Int, Int) !Int + deriving (Show) + +genQuery :: Int -> QC.Gen Query +genQuery n = do + QC.oneof + [ Rank0 <$> QC.chooseInt (0, n), + Rank1 <$> QC.chooseInt (0, n), + Select0 <$> QC.chooseInt (0, n), + Select1 <$> QC.chooseInt (0, n), + SelectKth0 <$> intervalGen n <*> QC.chooseInt (0, n), + SelectKth1 <$> intervalGen n <*> QC.chooseInt (0, n) + ] + +-- | Arbitrary return type for the `Query` result. +data Result + = I {-# UNPACK #-} !Int + | M !(Maybe Int) + deriving (Show, Eq) + +-- | containers. (referencial implementation) +handleRef :: VU.Vector Bool -> Query -> Result +handleRef xs q = case q of + Rank0 k -> I . VU.length . VU.filter not $ VU.take k xs + Rank1 k -> I . VU.length . VU.filter id $ VU.take k xs + Select0 k -> M . (fst <$>) . (VU.!? k) . VU.filter (not . snd) $ VU.indexed xs + Select1 k -> M . (fst <$>) . (VU.!? k) . VU.filter snd $ VU.indexed xs + SelectKth0 (!l, !r) k -> M . (fst <$>) . (VU.!? k) . VU.filter (not . snd) . VU.take (r - l) . VU.drop l $ VU.indexed xs + SelectKth1 (!l, !r) k -> M . (fst <$>) . (VU.!? k) . VU.filter snd . VU.take (r - l) . VU.drop l $ VU.indexed xs + +handleAcl :: BV.BitVector -> Query -> Result +handleAcl bv q = case q of + Rank0 k -> I $ BV.rank0 bv k + Rank1 k -> I $ BV.rank1 bv k + Select0 k -> M $ BV.select0 bv k + Select1 k -> M $ BV.select1 bv k + SelectKth0 (!l, !r) k -> M $ BV.selectKthIn0 bv l r k + SelectKth1 (!l, !r) k -> M $ BV.selectKthIn1 bv l r k + +prop_randomTest :: Init -> QC.Gen QC.Property +prop_randomTest Init {..} = do + qs <- QC.vectorOf capacity (genQuery capacity) + pure . QC.conjoin $ + map + ( \q -> + QC.counterexample (show q) $ + handleRef ref q QC.=== handleAcl bv q + ) + qs + +tests :: [TestTree] +tests = + [ QC.testProperty "random test" prop_randomTest + ] diff --git a/test/Tests/Extra/WaveletMatrix/Raw.hs b/test/Tests/Extra/WaveletMatrix/Raw.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/WaveletMatrix/Raw.hs @@ -0,0 +1,271 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.WaveletMatrix.Raw (tests) where + +import AtCoder.Extra.Bisect (lowerBound) +import AtCoder.Extra.WaveletMatrix.Raw qualified as WM +import Control.Exception (evaluate) +import Data.IntMap qualified as IM +import Data.Maybe (fromJust) +import Data.Ord (comparing) +import Data.Vector.Algorithms.Intro qualified as VAI +import Data.Vector.Unboxed qualified as VU +import System.IO.Unsafe (unsafePerformIO) +import Test.Hspec +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.Hspec +import Test.Tasty.QuickCheck as QC +import Tests.Util (intervalGen) + +compress :: VU.Vector Int -> VU.Vector Int +compress xs = + let dict = VU.uniq $ VU.modify VAI.sort xs + in VU.map (fromJust . lowerBound dict) xs + +data Init = Init + { capacity :: {-# UNPACK #-} !Int, + xs :: !(VU.Vector Int), + wm :: !WM.RawWaveletMatrix + } + deriving (Eq, Show) + +instance QC.Arbitrary Init where + arbitrary = do + QC.NonNegative n <- QC.arbitrary + xs <- compress . VU.fromListN n <$> QC.vectorOf n (QC.arbitrary @Int) + pure $ Init n xs (WM.build n xs) + +data Query + = Access !Int + | KthSmallestIn !(Int, Int) !Int + | IKthSmallestIn !(Int, Int) !Int + | KthLargestIn !(Int, Int) !Int + | IKthLargestIn !(Int, Int) !Int + | UnsafeKthSmallestIn !(Int, Int) !Int + | UnsafeIKthSmallestIn !(Int, Int) !Int + | UnsafeKthLargestIn !(Int, Int) !Int + | UnsafeIKthLargestIn !(Int, Int) !Int + | RankLT !(Int, Int) !Int + | RankBetween !(Int, Int) !(Int, Int) + | Rank !(Int, Int) !Int + | Select !Int + | SelectKth !Int !Int + | SelectKthIn !(Int, Int) !Int !Int + | LookupLE !(Int, Int) !Int + | LookupLT !(Int, Int) !Int + | LookupGE !(Int, Int) !Int + | LookupGT !(Int, Int) !Int + | AssocsIn !(Int, Int) + | DescAssocsIn !(Int, Int) + deriving (Show) + +genQuery :: Int -> QC.Gen Query +genQuery n = do + QC.oneof + [ Access <$> QC.chooseInt (-1, n), + KthSmallestIn <$> lr <*> exc, + IKthSmallestIn <$> lr <*> exc, + KthLargestIn <$> lr <*> exc, + IKthLargestIn <$> lr <*> exc, + UnsafeKthSmallestIn <$> lr <*> inc, + UnsafeIKthSmallestIn <$> lr <*> inc, + UnsafeKthLargestIn <$> lr <*> inc, + UnsafeIKthLargestIn <$> lr <*> inc, + RankLT <$> lr <*> val, + RankBetween <$> lr <*> lr, + Rank <$> lr <*> val, + Select <$> val, + SelectKth <$> exc <*> val, + SelectKthIn <$> lr <*> exc <*> val, + LookupLE <$> lr <*> QC.chooseInt (-1, n), + LookupLT <$> lr <*> QC.chooseInt (-1, n), + LookupGE <$> lr <*> QC.chooseInt (-1, n), + LookupGT <$> lr <*> QC.chooseInt (-1, n), + AssocsIn <$> lr, + DescAssocsIn <$> lr + ] + where + inc = QC.chooseInt (0, n - 1) + exc = QC.chooseInt (0, n) + lr = intervalGen n + val = QC.chooseInt (-n, n) + +-- | Arbitrary return type for the `Query` result. +data Result + = I {-# UNPACK #-} !Int + | M !(Maybe Int) + | M2 !(Maybe (Int, Int)) + | Assocs [(Int, Int)] + deriving (Show, Eq) + +-- | containers. (referencial implementation) +handleRef :: VU.Vector Int -> Query -> Result +handleRef xs q = case q of + Access i -> M $ xs VU.!? i + KthSmallestIn (!l, !r) k -> M $ snd <$> ikthSmallestIn l r k + IKthSmallestIn (!l, !r) k -> M2 $ ikthSmallestIn l r k + KthLargestIn (!l, !r) k -> M $ snd <$> ikthLargestIn l r k + IKthLargestIn (!l, !r) k -> M2 $ ikthLargestIn l r k + UnsafeKthSmallestIn (!l, !r) k -> M $ snd <$> ikthSmallestIn l r k + UnsafeIKthSmallestIn (!l, !r) k -> M2 $ ikthSmallestIn l r k + UnsafeKthLargestIn (!l, !r) k -> M $ snd <$> ikthLargestIn l r k + UnsafeIKthLargestIn (!l, !r) k -> M2 $ ikthLargestIn l r k + RankLT (!l, !r) xr -> rankBetween l r (minBound `div` 2) xr + RankBetween (!l, !r) (!xl, !xr) -> rankBetween l r xl xr + Rank (!l, !r) x -> rankBetween l r x (x + 1) + Select x -> M $ selectKthIn 0 n 0 x + SelectKth k x -> M $ selectKthIn 0 n k x + SelectKthIn (!l, !r) k x -> M $ selectKthIn l r k x + LookupLE (!l, !r) x -> max_ . VU.filter (<= x) . VU.take (r - l) $ VU.drop l xs + LookupLT (!l, !r) x -> max_ . VU.filter (< x) . VU.take (r - l) $ VU.drop l xs + LookupGE (!l, !r) x -> min_ . VU.filter (>= x) . VU.take (r - l) $ VU.drop l xs + LookupGT (!l, !r) x -> min_ . VU.filter (> x) . VU.take (r - l) $ VU.drop l xs + AssocsIn (!l, !r) -> Assocs . IM.assocs . IM.fromListWith (+) . VU.toList . VU.map (,1) . VU.take (r - l) $ VU.drop l xs + DescAssocsIn (!l, !r) -> Assocs . reverse . IM.assocs . IM.fromListWith (+) . VU.toList . VU.map (,1) . VU.take (r - l) $ VU.drop l xs + where + n = VU.length xs + max_ ys + | VU.null ys = M Nothing + | otherwise = M $ Just $ VU.maximum ys + min_ ys + | VU.null ys = M Nothing + | otherwise = M $ Just $ VU.minimum ys + ikthSmallestIn l r k = + (VU.!? k) + . VU.modify (VAI.sortBy (comparing (\(!i, !x) -> (x, i)))) + . VU.take (r - l) + . VU.drop l + . VU.indexed + $ xs + ikthLargestIn l r k = ikthSmallestIn l r ((r - l) - (k + 1)) + rankBetween l r xl xr = + I + . VU.length + . VU.filter (\x -> xl <= x && x < xr) + . VU.take (r - l) + . VU.drop l + $ xs + selectKthIn l r k x = + (fst <$>) + . (VU.!? k) + . VU.filter ((== x) . snd) + . VU.take (r - l) + . VU.drop l + . VU.indexed + $ xs + +handleAcl :: WM.RawWaveletMatrix -> Query -> Result +handleAcl wm q = case q of + Access i -> M $ WM.access wm i + KthSmallestIn (!l, !r) k -> M $ WM.kthSmallestIn wm l r k + IKthSmallestIn (!l, !r) k -> M2 $ WM.ikthSmallestIn wm l r k + KthLargestIn (!l, !r) k -> M $ WM.kthLargestIn wm l r k + IKthLargestIn (!l, !r) k -> M2 $ WM.ikthLargestIn wm l r k + UnsafeKthSmallestIn (!l, !r) k -> M $ WM.kthSmallestIn wm l r k + UnsafeIKthSmallestIn (!l, !r) k -> M2 $ WM.ikthSmallestIn wm l r k + UnsafeKthLargestIn (!l, !r) k -> M $ WM.kthLargestIn wm l r k + UnsafeIKthLargestIn (!l, !r) k -> M2 $ WM.ikthLargestIn wm l r k + RankLT (!l, !r) xr -> I $ WM.rankLT wm l r xr + RankBetween (!l, !r) (!xl, !xr) -> I $ WM.rankBetween wm l r xl xr + Rank (!l, !r) x -> I $ WM.rank wm l r x + Select x -> M $ WM.select wm x + SelectKth k x -> M $ WM.selectKth wm k x + SelectKthIn (!l, !r) k x -> M $ WM.selectKthIn wm l r k x + LookupLE (!l, !r) x -> M $ WM.lookupLE wm l r x + LookupLT (!l, !r) x -> M $ WM.lookupLT wm l r x + LookupGE (!l, !r) x -> M $ WM.lookupGE wm l r x + LookupGT (!l, !r) x -> M $ WM.lookupGT wm l r x + AssocsIn (!l, !r) -> Assocs $ WM.assocsIn wm l r + DescAssocsIn (!l, !r) -> Assocs $ WM.descAssocsIn wm l r + +prop_randomTest :: Init -> QC.Gen QC.Property +prop_randomTest Init {..} = do + qs <- QC.vectorOf capacity (genQuery capacity) + pure . QC.conjoin $ + map + ( \q -> + QC.counterexample (show q) $ + handleRef xs q QC.=== handleAcl wm q + ) + qs + +unit_boundary :: TestTree +unit_boundary = testCase "boundary" $ do + let n = 5 + let wm = WM.build 3 $ VU.fromList [0, 1, 2, 1, 0] + + let try :: (HasCallStack, Eq a, Show a) => (WM.RawWaveletMatrix -> Int -> Int -> Int -> Maybe a) -> Int -> IO () + try f x = do + (@?= Nothing) $ f wm (-1) 0 x + (@?= Nothing) $ f wm n (n + 1) x + + let k = 0 + try WM.kthSmallestIn k + try WM.ikthSmallestIn k + try WM.kthLargestIn k + try WM.ikthLargestIn k + +-- try WM.unsafeKthSmallestIn +-- try WM.unsafeIKthSmallestIn +-- try WM.unsafeKthLargestIn +-- try WM.unsafeIKthLargestIn + + let tryRank :: (HasCallStack) => (WM.RawWaveletMatrix -> Int -> Int -> Int -> Int) -> Int -> IO () + tryRank f x = do + -- out of range + (@?= 0) $ f wm (-1) 0 x + (@?= 0) $ f wm n (n + 1) x + -- reverse + (@?= 0) $ f wm 1 0 x + (@?= 0) $ f wm n (n - 1) x + (@?= 0) $ f wm n 0 x + -- out of range and reverse + (@?= 0) $ f wm 0 (-1) x + (@?= 0) $ f wm (n + 1) n x + + let tryRankBetween :: (HasCallStack) => (WM.RawWaveletMatrix -> Int -> Int -> Int -> Int -> Int) -> Int -> Int -> IO () + tryRankBetween f xl xr = do + -- out of range + (@?= 0) $ f wm (-1) 0 xl xr + (@?= 0) $ f wm n (n + 1) xl xr + -- reverse + (@?= 0) $ f wm 1 0 xl xr + (@?= 0) $ f wm n (n - 1) xl xr + (@?= 0) $ f wm n 0 xl xr + -- out of range and reverse + (@?= 0) $ f wm 0 (-1) xl xr + (@?= 0) $ f wm (n + 1) n xl xr + + tryRank WM.rankLT 0 + tryRank WM.rankLT 1 + tryRank WM.rankLT (-1) + tryRankBetween WM.rankBetween (-1) 1 + tryRank WM.rank 0 + tryRank WM.rank 1 + tryRank WM.rank (-1) + +-- TODO: test + +-- (@?= Nothing) $ WM.select wm +-- (@?= Nothing) $ WM.selectKth wm +-- (@?= Nothing) $ WM.selectKthIn wm + +-- (@?= Nothing) $ WM.lookupLE wm +-- (@?= Nothing) $ WM.lookupLT wm +-- (@?= Nothing) $ WM.lookupGE wm +-- (@?= Nothing) $ WM.lookupGT wm +-- (@?= Nothing) $ WM.assocsIn wm +-- (@?= Nothing) $ WM.descAssocsIn wm + +spec_invalid :: IO TestTree +spec_invalid = testSpec "boundary check" $ do + it "throws error 1" $ do + evaluate (WM.build (-1) (VU.singleton 1)) `shouldThrow` anyException + +tests :: [TestTree] +tests = + [ unit_boundary, + unsafePerformIO spec_invalid, + QC.testProperty "random test" prop_randomTest + ] diff --git a/test/Tests/Extra/WaveletMatrix2d.hs b/test/Tests/Extra/WaveletMatrix2d.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/WaveletMatrix2d.hs @@ -0,0 +1,119 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.WaveletMatrix2d (tests) where + +import AtCoder.Extra.WaveletMatrix2d qualified as WM +import Control.Monad (foldM_) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (RealWorld) +import Data.List qualified as L +import Data.Map qualified as M +import Data.Maybe (fromMaybe) +import Data.Ord (comparing) +import Data.Semigroup (Sum (..)) +import Data.Vector.Algorithms.Intro qualified as VAI +import Data.Vector.Unboxed qualified as VU +import Test.Hspec +import Test.QuickCheck.Monadic as QCM +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck as QC +import Tests.Util (intervalGen) + +data Init = Init + { capacity :: {-# UNPACK #-} !Int, + map0 :: !(M.Map (Int, Int) (Sum Int)), + wmM :: !(IO (WM.WaveletMatrix2d RealWorld (Sum Int))) + } + +instance Show Init where + show Init {..} = show ("Init", capacity, map0) + +instance QC.Arbitrary Init where + arbitrary = do + QC.NonNegative n <- QC.arbitrary + let yxs = VU.fromList [(x, y) | x <- [-16 .. 16], y <- [-16 .. 16]] + pure $ Init n M.empty (WM.new negate yxs) + +data Query + = Read !(Int, Int) + | Write !(Int, Int) !Int + | ModifyAdd !Int !(Int, Int) + | Prod !(Int, Int) !(Int, Int) + | ProdMaybe !(Int, Int) !(Int, Int) + | AllProd + deriving (Show) + +genQuery :: Int -> QC.Gen Query +genQuery n = do + QC.oneof + [ Read <$> lr, + Write <$> lr <*> val, + ModifyAdd <$> val <*> lr, + Prod <$> lr <*> lr, + ProdMaybe <$> lr <*> lr, + pure AllProd + ] + where + lr = (\(!x, !y) -> (x - 16, y - 16)) <$> intervalGen 32 + val = QC.arbitrary @Int + +-- | Arbitrary return type for the `Query` result. +data Result + = None + | S !(Sum Int) + | MS !(Maybe (Sum Int)) + deriving (Show, Eq) + +-- | containers. (referencial implementation) +handleRef :: M.Map (Int, Int) (Sum Int) -> Query -> (Result, M.Map (Int, Int) (Sum Int)) +handleRef map q = case q of + Read (!x, !y) -> (S . fromMaybe mempty $ M.lookup (x, y) map, map) + Write (!x, !y) v -> (None, M.insert (x, y) (Sum v) map) + ModifyAdd w (!x, !y) -> (None, M.insertWith (+) (x, y) (Sum w) map) + Prod (!x1, !x2) (!y1, !y2) -> (S $ prod x1 x2 y1 y2, map) + ProdMaybe (!x1, !x2) (!y1, !y2) -> (MS . Just $ prod x1 x2 y1 y2, map) + AllProd -> (S $ L.foldl' (<>) mempty (M.elems map), map) + where + prod x1 x2 y1 y2 = + L.foldl' (<>) (mempty :: Sum Int) + . (snd <$>) + . filter (\((!x, !y), !_) -> x1 <= x && x < x2 && y1 <= y && y < y2) + $ M.assocs map + +handleAcl :: (PrimMonad m) => WM.WaveletMatrix2d (PrimState m) (Sum Int) -> Query -> m Result +handleAcl wm q = case q of + Read (!x, !y) -> do + S <$> WM.read wm (x, y) + Write (!x, !y) v -> do + WM.write wm (x, y) $ Sum v + pure None + ModifyAdd w (!x, !y) -> do + WM.modify wm (+ Sum w) (x, y) + pure None + Prod (!x1, !x2) (!y1, !y2) -> do + S <$> WM.prod wm x1 x2 y1 y2 + ProdMaybe (!x1, !x2) (!y1, !y2) -> do + MS <$> WM.prodMaybe wm x1 x2 y1 y2 + AllProd -> do + S <$> WM.allProd wm + +prop_randomTest :: Init -> QC.Property +prop_randomTest Init {..} = QCM.monadicIO $ do + wm <- QCM.run wmM + qs <- QCM.pick $ QC.vectorOf capacity (genQuery capacity) + foldM_ + ( \map query -> do + let (!expected, !map') = handleRef map query + actual <- QCM.run $ handleAcl wm query + QCM.assertWith (expected == actual) $ show (query, expected, actual) + pure map' + ) + map0 + qs + +tests :: [TestTree] +tests = + [ -- unit_boundary, + QC.testProperty "random test" prop_randomTest + ] diff --git a/test/Tests/MaxFlow.hs b/test/Tests/MaxFlow.hs --- a/test/Tests/MaxFlow.hs +++ b/test/Tests/MaxFlow.hs @@ -2,7 +2,7 @@ module Tests.MaxFlow (tests) where import AtCoder.MaxFlow qualified as MF -import Data.Bit (Bit(..)) +import Data.Bit (Bit (..)) import Data.Vector.Unboxed qualified as VU import System.IO.Unsafe (unsafePerformIO) import Test.Hspec diff --git a/test/Tests/Scc.hs b/test/Tests/Scc.hs --- a/test/Tests/Scc.hs +++ b/test/Tests/Scc.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE RecordWildCards #-} - module Tests.Scc (tests) where import AtCoder.Scc qualified as Scc diff --git a/test/Tests/Util.hs b/test/Tests/Util.hs --- a/test/Tests/Util.hs +++ b/test/Tests/Util.hs @@ -1,4 +1,4 @@ -module Tests.Util (myForAllShrink, laws) where +module Tests.Util (myForAllShrink, laws, intervalGen) where import Data.Proxy (Proxy (..)) import Data.Typeable (Typeable, typeRep) @@ -40,3 +40,10 @@ let QCC.Laws name pairs = f (Proxy @a) in testGroup name (map (uncurry QC.testProperty) pairs) ) + +-- | Returns an interval [l, r) in [0, n) +intervalGen :: Int -> QC.Gen (Int, Int) +intervalGen n = do + l <- QC.chooseInt (0, n) + r <- QC.chooseInt (l, n) + pure (l, r)