diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,14 @@ # Revision history for acl-hs +## 1.2.0.0 -- Feb 2025 + +- Added `AtCoder.Extra.Seq` +- Tweaked `INLINE` settings for less compile time +- Breaking changes: + - `Matrix.diag` now does not take length parameter + - `Extra.Math.primitiveRoot` is renamed to `primitiveRoot32` + - `Internal.Convolution` functions now use `ST` instead of `PrimMonad` + ## 1.1.1.0 -- Jan 2025 - Added `AtCoder.Extra.Tree.Lct` diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ - The library is mainly for AtCoder and only GHC 9.8.4 is guaranteed to be supported. - Functions primarily use half-open interval [l, r). -- The extra module contains additional utilities beyond the original C++ library. +- 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,14 +4,14 @@ -- PVP summary: +-+------- breaking API changes -- | | +----- non-breaking API additions -- | | | +--- code changes with no API change -version: 1.1.1.0 +version: 1.2.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. + - The `Extra` module contains additional utilities beyond the original C++ library. category: Algorithms, Data Structures license: CC0-1.0 @@ -34,17 +34,6 @@ ghc-options: -Wall common dependencies - -- AtCoder environment (2023 -) - -- if impl(ghc ==9.4.5) - -- build-depends: - -- , base ==4.17.1.0 - -- , bitvec ^>=1.1.4.0 - -- , bytestring ^>=0.11.4.0 - -- , primitive ^>=0.8.0.0 - -- , vector ^>=0.13.0.0 - -- , vector-algorithms ^>=0.9.0.1 - -- , wide-word - build-depends: , base >=4.9 && <4.22 , bitvec <1.2 @@ -89,8 +78,11 @@ AtCoder.Extra.Monoid.V2 AtCoder.Extra.MultiSet AtCoder.Extra.Pdsu + AtCoder.Extra.Pool AtCoder.Extra.Semigroup.Matrix AtCoder.Extra.Semigroup.Permutation + AtCoder.Extra.Seq + AtCoder.Extra.Seq.Raw AtCoder.Extra.Tree AtCoder.Extra.Tree.Hld AtCoder.Extra.Tree.Lct @@ -151,6 +143,7 @@ Tests.Extra.MultiSet Tests.Extra.Semigroup.Matrix Tests.Extra.Semigroup.Permutation + Tests.Extra.Seq Tests.Extra.WaveletMatrix Tests.Extra.WaveletMatrix.BitVector Tests.Extra.WaveletMatrix.Raw diff --git a/src/AtCoder/Convolution.hs b/src/AtCoder/Convolution.hs --- a/src/AtCoder/Convolution.hs +++ b/src/AtCoder/Convolution.hs @@ -76,7 +76,7 @@ -- - \(O(n\log{n} + \log{\mathrm{mod}})\), where \(n = |a| + |b|\). -- -- @since 1.0.0.0 -{-# INLINABLE convolution #-} +{-# INLINE convolution #-} convolution :: forall p. (HasCallStack, AM.Modulus p) => @@ -107,7 +107,7 @@ -- - \(O(n\log{n} + \log{\mathrm{mod}})\), where \(n = |a| + |b|\). -- -- @since 1.0.0.0 -{-# INLINABLE convolutionRaw #-} +{-# INLINE convolutionRaw #-} convolutionRaw :: forall p a. (HasCallStack, AM.Modulus p, Integral a, VU.Unbox a) => @@ -137,7 +137,7 @@ -- - \(O(n\log{n})\), where \(n = |a| + |b|\). -- -- @since 1.0.0.0 -{-# INLINABLE convolution64 #-} +{-# INLINE convolution64 #-} convolution64 :: (HasCallStack) => VU.Vector Int -> diff --git a/src/AtCoder/Dsu.hs b/src/AtCoder/Dsu.hs --- a/src/AtCoder/Dsu.hs +++ b/src/AtCoder/Dsu.hs @@ -65,7 +65,8 @@ import AtCoder.Internal.Assert qualified as ACIA import Control.Monad (when) -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (ST) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) import Data.Vector qualified as V import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM @@ -117,11 +118,11 @@ -- @since 1.0.0.0 {-# INLINE merge #-} merge :: (HasCallStack, PrimMonad m) => Dsu (PrimState m) -> Int -> Int -> m Int -merge dsu@Dsu {..} a b = do +merge dsu@Dsu {..} a b = stToPrim $ do let !_ = ACIA.checkVertex "AtCoder.Dsu.merge" a nDsu let !_ = ACIA.checkVertex "AtCoder.Dsu.merge" b nDsu - x <- leader dsu a - y <- leader dsu b + x <- leaderST dsu a + y <- leaderST dsu b if x == y then do pure x @@ -169,6 +170,17 @@ lb <- leader dsu b pure $ la == lb +{-# INLINE leaderST #-} +leaderST :: Dsu s -> Int -> ST s Int +leaderST dsu@Dsu {..} a = do + pa <- VGM.read parentOrSizeDsu a + if pa < 0 + then pure a + else do + lpa <- leaderST dsu pa + VGM.write parentOrSizeDsu a lpa + pure lpa + -- | Returns the representative of the connected component that contains the vertex \(a\). -- -- ==== Constraints @@ -180,15 +192,7 @@ -- @since 1.0.0.0 {-# INLINE leader #-} leader :: (HasCallStack, PrimMonad m) => Dsu (PrimState m) -> Int -> m Int -leader dsu@Dsu {..} a = do - let !_ = ACIA.checkVertex "AtCoder.Dsu.leader" a nDsu - pa <- VGM.read parentOrSizeDsu a - if pa < 0 - then pure a - else do - lpa <- leader dsu pa - VGM.write parentOrSizeDsu a lpa - pure lpa +leader dsu a = stToPrim $ leaderST dsu a -- | Returns the size of the connected component that contains the vertex \(a\). -- @@ -201,9 +205,9 @@ -- @since 1.0.0.0 {-# INLINE size #-} size :: (HasCallStack, PrimMonad m) => Dsu (PrimState m) -> Int -> m Int -size dsu@Dsu {..} a = do +size dsu@Dsu {..} a = stToPrim $ do let !_ = ACIA.checkVertex "AtCoder.Dsu.size" a nDsu - la <- leader dsu a + la <- leaderST dsu a sizeLa <- VGM.read parentOrSizeDsu la pure (-sizeLa) @@ -218,10 +222,10 @@ -- @since 1.0.0.0 {-# INLINE groups #-} groups :: (PrimMonad m) => Dsu (PrimState m) -> m (V.Vector (VU.Vector Int)) -groups dsu@Dsu {..} = do +groups dsu@Dsu {..} = stToPrim $ do groupSize <- VUM.replicate nDsu (0 :: Int) leaders <- VU.generateM nDsu $ \i -> do - li <- leader dsu i + li <- leaderST dsu i VGM.modify groupSize (+ 1) li pure li result <- do diff --git a/src/AtCoder/Extra/HashMap.hs b/src/AtCoder/Extra/HashMap.hs --- a/src/AtCoder/Extra/HashMap.hs +++ b/src/AtCoder/Extra/HashMap.hs @@ -6,8 +6,10 @@ -- | 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/). +-- ==== Capacity limitation +-- Access to each key creates a new entry. Note that entries cannot be invalidated due to the +-- internal implementation (called /open addressing/). If the hash map is full, +-- __access to a new key causes an infinite loop__ . -- -- ==== __Example__ -- Create a `HashMap` with `capacity` \(10\): @@ -76,7 +78,8 @@ import AtCoder.Internal.Assert qualified as ACIA import Control.Monad (void, when) -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (ST) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) import Data.Bit (Bit (..)) import Data.Bits (Bits (xor, (.&.)), (.>>.)) import Data.Vector.Generic qualified as VG @@ -171,10 +174,14 @@ -- | \(O(1)\) (Internal) Hashed slot search. -- +-- ==== Constraint +-- - The rest capacity must be non-zero. Otherwise it loops forever. +-- -- @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) +{-# INLINE indexST #-} +indexST :: (HasCallStack) => HashMap s a -> Int -> ST s Int +indexST hm@HashMap {..} k = do + inner (hash hm k) where inner !h = do Bit b <- VGM.read usedHM h @@ -187,11 +194,14 @@ -- | \(O(1)\) Return the value to which the specified key is mapped, or `Nothing` if this map -- contains no mapping for the key. -- +-- ==== Constraint +-- - The rest capacity must be non-zero. Otherwise it loops forever. +-- -- @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 + i <- stToPrim $ indexST hm k Bit b <- VGM.read usedHM i if b then Just <$> VGM.read valHM i @@ -199,11 +209,14 @@ -- | \(O(1)\) Checks whether the hash map contains the element. -- +-- ==== Constraint +-- - The rest capacity must be non-zero. Otherwise it loops forever. +-- -- @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 + i <- stToPrim $ indexST hm k Bit b <- VGM.read usedHM i -- TODO: is this key check necessary k' <- VGM.read keyHM i @@ -211,6 +224,9 @@ -- | \(O(1)\) Checks whether the hash map does not contain the element. -- +-- ==== Constraint +-- - The rest capacity must be non-zero. Otherwise it loops forever. +-- -- @since 1.1.0.0 {-# INLINE notMember #-} notMember :: (HasCallStack, PrimMonad m) => HashMap (PrimState m) a -> Int -> m Bool @@ -218,6 +234,9 @@ -- | \(O(1)\) Inserts a \((k, v)\) pair. -- +-- ==== Constraint +-- - The rest capacity must be non-zero. Otherwise it loops forever. +-- -- @since 1.1.0.0 {-# INLINE insert #-} insert :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> Int -> a -> m () @@ -226,11 +245,14 @@ -- | \(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}}))\). -- +-- ==== Constraint +-- - The rest capacity must be non-zero. Otherwise it loops forever. +-- -- @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 + i <- stToPrim $ indexST hm k Bit b <- VGM.exchange usedHM i $ Bit True if b then do @@ -245,11 +267,14 @@ -- | \(O(1)\) Inserts a \((k, v)\) pair and returns the old value, or `Nothing` if no such entry -- exists. -- +-- ==== Constraint +-- - The rest capacity must be non-zero. Otherwise it loops forever. +-- -- @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 + i <- stToPrim $ indexST hm k Bit b <- VGM.exchange usedHM i $ Bit True if b then do @@ -268,7 +293,7 @@ {-# 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 + i <- stToPrim $ indexST hm k Bit b <- VGM.read usedHM i when b $ do VGM.modify valHM f i @@ -279,7 +304,7 @@ {-# 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 + i <- stToPrim $ indexST hm k Bit b <- VGM.read usedHM i when b $ do VGM.modifyM valHM f i diff --git a/src/AtCoder/Extra/IntervalMap.hs b/src/AtCoder/Extra/IntervalMap.hs --- a/src/AtCoder/Extra/IntervalMap.hs +++ b/src/AtCoder/Extra/IntervalMap.hs @@ -64,7 +64,7 @@ -- * Lookups contains, - intersects, + containsInterval, lookup, read, readMaybe, @@ -164,15 +164,15 @@ -- @since 1.1.0.0 {-# INLINE contains #-} contains :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> m Bool -contains itm i = intersects itm i (i + 1) +contains itm i = containsInterval 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 -{-# INLINE intersects #-} -intersects :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m Bool -intersects (IntervalMap dim) l r +{-# INLINE containsInterval #-} +containsInterval :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m Bool +containsInterval (IntervalMap dim) l r | l >= r = pure False | otherwise = do res <- IM.lookupLE dim l 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 @@ -5,7 +5,7 @@ ( -- * Re-exports from the internal math module isPrime32, ACIM.invGcd, - ACIM.primitiveRoot, + primitiveRoot32, -- * Binary exponentiation @@ -26,8 +26,10 @@ ) where +import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Math qualified as ACIM import Data.Bits ((.>>.)) +import GHC.Stack (HasCallStack) -- | \(O(k \log^3 n) (k = 3)\). Returns whether the given `Int` value is a prime number. -- @@ -38,8 +40,23 @@ -- -- @since 1.1.0.0 {-# INLINE isPrime32 #-} -isPrime32 :: Int -> Bool -isPrime32 = ACIM.isPrime +isPrime32 :: (HasCallStack) => Int -> Bool +isPrime32 x = ACIM.isPrime x + where + !_ = ACIA.runtimeAssert (x < 4759123141) $ "AtCoder.Extra.Math.isPrime32: given too large number `" ++ show x ++ "`" + +-- | Returns the primitive root of the given `Int`. +-- +-- ==== Constraints +-- - The input must be a prime number. +-- - The input must be less than \(2^32\). +-- +-- @since 1.2.0.0 +{-# INLINE primitiveRoot32 #-} +primitiveRoot32 :: (HasCallStack) => Int -> Int +primitiveRoot32 x = ACIM.primitiveRoot x + where + !_ = ACIA.runtimeAssert (x < (1 .>>. 32)) $ "AtCoder.Extra.Math.primitiveRoot32: given too large number `" ++ show x ++ "`" -- | Calculates \(x^n\) with custom multiplication operator using the binary exponentiation -- technique. 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 @@ -17,7 +17,7 @@ where import AtCoder.LazySegTree (SegAct (..)) -import Data.Semigroup (stimes, Sum (..), Max(..), Min(..)) +import Data.Semigroup (Max (..), Min (..), Sum (..), stimes) import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU @@ -25,15 +25,24 @@ -- | Monoid action \(f: x \rightarrow x + d\). -- --- ==== __Example__ +-- ==== __Example (action on @Sum@)__ -- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeAdd(..)) -- >>> import AtCoder.LazySegTree qualified as LST --- >>> import Data.Semigroup (Max(..)) +-- >>> import Data.Semigroup (Sum(..)) -- >>> 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 -- +-- ==== __Example (action on @Max@)__ +-- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeAdd(..)) +-- >>> import AtCoder.LazySegTree qualified as LST +-- >>> import Data.Semigroup (Max(..)) +-- >>> seg <- LST.build @_ @(RangeAdd (Max Int)) @(Max Int) $ VU.generate 3 Max -- [0, 1, 2] +-- >>> LST.applyIn seg 0 3 $ RangeAdd (Max 5) -- [5, 6, 7] +-- >>> getMax <$> LST.prod seg 0 3 +-- 7 +-- -- @since 1.0.0.0 newtype RangeAdd a = RangeAdd a deriving newtype @@ -66,37 +75,38 @@ act :: (Semigroup a) => RangeAdd a -> a -> a act (RangeAdd dx) x = dx <> x --- | \(O(1)\) Acts on @a@ with length in terms of `SegAct`. +-- | \(O(1)\) Acts on @a@ with length in terms of `SegAct`. It doesn't work well with idempotent +-- monoids such as `Max` or `Min`. -- -- @since 1.0.0.0 {-# INLINE actWithLength #-} actWithLength :: (Semigroup a) => Int -> RangeAdd a -> a -> a actWithLength len (RangeAdd f) x = stimes len f <> x --- | @since 1.0.0.0 -instance (Semigroup a) => Semigroup (RangeAdd a) where +-- | @since 1.2.0.0 +instance (Semigroup a, Num a) => Semigroup (RangeAdd a) where {-# INLINE (<>) #-} - (RangeAdd a) <> (RangeAdd b) = RangeAdd $! a <> b + (RangeAdd a) <> (RangeAdd b) = RangeAdd $! a + b --- | @since 1.1.0.0 -instance (Monoid a) => Monoid (RangeAdd a) where +-- | @since 1.2.0.0 +instance (Num a, Semigroup a) => Monoid (RangeAdd a) where {-# INLINE mempty #-} - mempty = RangeAdd mempty + mempty = RangeAdd 0 --- | @since 1.1.0.0 -instance (Monoid (Sum a)) => SegAct (RangeAdd (Sum a)) (Sum a) where +-- | @since 1.2.0.0 +instance (Num a) => SegAct (RangeAdd (Sum a)) (Sum a) where {-# INLINE segActWithLength #-} - segActWithLength len f x = actWithLength len f x + segActWithLength = actWithLength -- | @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 +instance (Num a, Monoid (Max a)) => SegAct (RangeAdd (Max a)) (Max a) where + {-# INLINE segAct #-} + segAct (RangeAdd (Max dx)) (Max x) = Max $! dx + 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 +instance (Num a, Monoid (Min a)) => SegAct (RangeAdd (Min a)) (Min a) where + {-# INLINE segAct #-} + segAct (RangeAdd (Min dx)) (Min x) = Min $! dx + 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/MultiSet.hs b/src/AtCoder/Extra/MultiSet.hs --- a/src/AtCoder/Extra/MultiSet.hs +++ b/src/AtCoder/Extra/MultiSet.hs @@ -3,13 +3,14 @@ -- | A fast, mutable multiset for `Int` keys backed by a @HashMap@. Most operations are performed -- in \(O(1)\) time, but in average. -- +-- ==== Capacity limitation +-- Access to each key creates a new entry. Note that entries cannot be invalidated due to the +-- internal implementation (called /open addressing/). If the hash map is full, +-- __access to a new key causes infinite loop__ . +-- -- ==== 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@. +-- violated on `add` or `sub`. -- -- ==== __Example__ -- Create a `MultiSet` with capacity \(4\): diff --git a/src/AtCoder/Extra/Pool.hs b/src/AtCoder/Extra/Pool.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Pool.hs @@ -0,0 +1,169 @@ +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Fixed-sized array for \(O(1)\) allocation and \(O(1)\) clearing after \(O(n)\) construction. +module AtCoder.Extra.Pool + ( -- * Pool + Pool (..), + Index (..), + undefIndex, + nullIndex, + + -- * Constructors + new, + clear, + + -- * Metadata + capacity, + size, + + -- * Allocations + alloc, + free, + + -- * Read/write + read, + write, + modify, + exchange, + ) +where + +import AtCoder.Internal.Buffer qualified as B +import Control.Monad.Primitive (PrimMonad, PrimState) +import Data.Coerce +import Data.Vector.Generic qualified as VG +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Primitive qualified as VP +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import Prelude hiding (read) + +-- | Fixed-sized array for \(O(1)\) allocation and \(O(1)\) clearing after \(O(n)\) construction. +data Pool s a = Pool + { -- | Data array. + dataPool :: !(VUM.MVector s a), + -- | Free slot indices pushed on free. + freePool :: !(B.Buffer s Index), + -- | Next index when `freePool` is empty. + nextPool :: !(VUM.MVector s Index) + } + +-- | Strongly typed index of pool items. User has to explicitly @corece@ on raw index use, but it's +-- ok as far as the end user don't see it. +newtype Index = Index {unIndex :: Int} + deriving (Eq, VP.Prim) + deriving newtype (Ord, Show) + +newtype instance VU.MVector s Index = MV_Index (VP.MVector s Index) + +newtype instance VU.Vector Index = V_Index (VP.Vector Index) + +deriving via (VU.UnboxViaPrim Index) instance VGM.MVector VUM.MVector Index + +deriving via (VU.UnboxViaPrim Index) instance VG.Vector VU.Vector Index + +instance VU.Unbox Index + +-- | Invalid, null `Index`. +{-# INLINE undefIndex #-} +undefIndex :: Index +undefIndex = Index (-1) + +-- | Returns `True` for `undefIndex`. +{-# INLINE nullIndex #-} +nullIndex :: Index -> Bool +nullIndex = (== undefIndex) + +-- | \(O(n)\) Creates a pool with the specified @capacity@. +{-# INLINE new #-} +new :: (VU.Unbox a, PrimMonad m) => Int -> m (Pool (PrimState m) a) +new capacity = do + dataPool <- VUM.unsafeNew capacity + freePool <- B.new capacity + nextPool <- VUM.replicate 1 (Index 0) + pure Pool {..} + +-- | \(O(1)\) Resets the pool to the initial state. +{-# INLINE clear #-} +clear :: (PrimMonad m) => Pool (PrimState m) a -> m () +clear Pool {..} = do + B.clear freePool + VGM.unsafeWrite nextPool 0 $ Index 0 + +-- | \(O(1)\) Returns the maximum number of elements the pool can store. +{-# INLINE capacity #-} +capacity :: (VU.Unbox a) => Pool s a -> Int +capacity = VGM.length . dataPool + +-- | \(O(1)\) Returns the number of elements in the pool. +{-# INLINE size #-} +size :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> m Int +size Pool {..} = do + !nFree <- B.length freePool + Index !next <- VGM.unsafeRead nextPool 0 + let !cap = VGM.length dataPool + pure $ cap - (next - nFree) + +-- | \(O(1)\) Allocates a new element. +-- +-- ==== Constraints +-- - The number of elements must not exceed the `capacity`. +{-# INLINE alloc #-} +alloc :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> a -> m Index +alloc Pool {..} !x = do + B.popBack freePool >>= \case + Just i -> pure i + Nothing -> do + Index i <- VGM.unsafeRead nextPool 0 + VGM.unsafeWrite nextPool 0 $ coerce (i + 1) + VGM.write dataPool i x + pure $ coerce i + +-- | \(O(1)\) Frees an element. Be sure to not free a deleted element. +-- +-- ==== Constraints +-- - \(0 \le i \lt n\) +{-# INLINE free #-} +free :: (PrimMonad m) => Pool (PrimState m) a -> Index -> m () +free Pool {..} i = do + B.pushBack freePool i + +-- | \(O(1)\) Reads the \(k\)-th value. +-- +-- ==== Constraints +-- - \(0 \le i \lt n\) +{-# INLINE read #-} +read :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> Index -> m a +read Pool {dataPool} !i = do + VGM.read dataPool (coerce i) + +-- | \(O(1)\) Writes to the \(k\)-th value. +-- +-- ==== Constraints +-- - \(0 \le i \lt n\) +{-# INLINE write #-} +write :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> Index -> a -> m () +write Pool {dataPool} !i !x = do + VGM.write dataPool (coerce i) x + +-- | \(O(1)\) Modifies the \(k\)-th value. +-- +-- ==== Constraints +-- - \(0 \le i \lt n\) +{-# INLINE modify #-} +modify :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> (a -> a) -> Index -> m () +modify Pool {dataPool} !f !i = do + VGM.modify dataPool f (coerce i) + +-- | \(O(1)\) Exchanges the \(k\)-th value. +-- +-- ==== Constraints +-- - \(0 \le i \lt n\) +{-# INLINE exchange #-} +exchange :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> Index -> a -> m a +exchange Pool {dataPool} !i !x = do + VGM.exchange dataPool (coerce i) x diff --git a/src/AtCoder/Extra/Semigroup/Matrix.hs b/src/AtCoder/Extra/Semigroup/Matrix.hs --- a/src/AtCoder/Extra/Semigroup/Matrix.hs +++ b/src/AtCoder/Extra/Semigroup/Matrix.hs @@ -125,18 +125,18 @@ VGM.write vec (i + n * i) 1 pure vec --- FIXME: diag should not take `n` - -- | \(O(n^2)\) Creates an NxN diagonal matrix. -- -- @since 1.1.0.0 {-# INLINE diag #-} -diag :: (VU.Unbox a, Num a) => Int -> VU.Vector a -> Matrix a -diag n xs = Matrix n n $ VU.create $ do +diag :: (VU.Unbox a, Num a) => VU.Vector a -> Matrix a +diag 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 + where + n = VU.length xs -- | \(O(n^2)\) Maps the `Matrix`. -- diff --git a/src/AtCoder/Extra/Seq.hs b/src/AtCoder/Extra/Seq.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Seq.hs @@ -0,0 +1,781 @@ +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Dynamic sequence of monoid values with monoid actions on them through the `SegAct` instance. +-- +-- ==== __Example__ +-- +-- Create a `Seq` storage of length \(10\): +-- +-- >>> import AtCoder.Extra.Monoid.RangeAdd qualified as RangeAdd +-- >>> import AtCoder.Extra.Seq qualified as Seq +-- >>> import AtCoder.LazySegTree (SegAct (..)) +-- >>> import Data.Semigroup (Sum (..)) +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> seq <- Seq.new @_ @(RangeAdd.RangeAdd (Sum Int)) @(Sum Int) 10 +-- +-- Allocate a sequence of \(0, 1, 2, 3\): +-- +-- >>> handle <- Seq.newSeq seq (VU.fromList [0, 1, 2, 3]) +-- +-- Get monoid products: +-- +-- >>> Seq.prodAll seq handle +-- Sum {getSum = 6} +-- +-- >>> Seq.prod seq handle 1 3 +-- Sum {getSum = 3} +-- +-- `read`, `write`, `modify` and `exchange` are available: +-- +-- >>> -- [0, 1, 2, 3] -> [0, 10, 2, 30] +-- >>> Seq.write seq handle 3 30 +-- >>> Seq.modify seq handle (* 10) 1 +-- +-- Actions can be performed with `SegAct` instances: +-- +-- >>> -- [0, 10, 2, 30] -> [0, 20, 12, 40] +-- >>> Seq.applyIn seq handle 1 4 (RangeAdd.new 10) +-- >>> Seq.prod seq handle 1 4 +-- Sum {getSum = 72} +-- +-- The sequence can be reversed if the action type is commutative: +-- +-- >>> Seq.reverse seq handle 0 4 +-- >>> VU.map getSum <$> Seq.freeze seq handle +-- [40,12,20,0] +-- +-- The sequence is dynamic and new elements can be inserted or deleted: +-- +-- >>> -- [40, 12, 20, 0] -> [40, 33, 12, 20, 0] +-- >>> Seq.insert seq handle 1 (Sum 33) +-- >>> -- [40, 33, 12, 20, 0] -> [40, 33, 12, 0] +-- >>> Seq.delete seq handle 3 +-- Sum {getSum = 20} +-- +-- >>> VU.map getSum <$> Seq.freeze seq handle +-- [40,33,12,0] +-- +-- The `Seq` storage can store multiple sequences: +-- +-- >>> handle2 <- Seq.newSeq seq (VU.generate 2 Sum) +-- >>> VU.map getSum <$> Seq.freeze seq handle2 +-- [0,1] +-- +-- Merge/split operations are available. `merge` functions mutate the given @handle@ to be the +-- merged sequence handle: +-- +-- >>> Seq.merge seq handle handle2 +-- >>> VU.map getSum <$> Seq.freeze seq handle +-- [40,33,12,0,0,1] +-- +-- `split` functions mutate the given @handle@ to be the leftmost one and returns right sequence +-- handles: +-- +-- >>> (handleM, handleR) <- Seq.split3 seq handle 2 4 +-- >>> VU.map getSum <$> Seq.freeze seq handle +-- [40,33] +-- +-- >>> VU.map getSum <$> Seq.freeze seq handleM +-- [12,0] +-- +-- >>> VU.map getSum <$> Seq.freeze seq handleR +-- [0,1] +-- +-- Bisection methods are available for monoid values and their products: +-- +-- >>> Seq.reset seq +-- >>> handle <- Seq.newSeq seq $ VU.generate 10 Sum +-- >>> Seq.ilowerBound seq handle (\_ x -> x < 5) +-- 5 +-- +-- >>> Seq.ilowerBoundProd seq handle (\_ x -> x < 5) +-- 3 +-- +-- @since 1.2.0.0 +module AtCoder.Extra.Seq + ( -- * Seq + Seq.Seq (..), + Handle (..), + newHandle, + nullHandle, + invalidateHandle, + + -- * Constructors + new, + reset, + free, + newNode, + newSeq, + + -- * Merge/split + merge, + merge3, + merge4, + split, + split3, + split4, + splitLr, + -- slice, -- because it returns a raw `P.Index`, use the `Raw.sliceST` instead + + -- * Read/write + read, + readMaybe, + write, + modify, + exchange, + + -- * Products + prod, + prodMaybe, + prodAll, + + -- * Applications + applyIn, + applyToRoot, + reverse, + + -- * Insert/delete + insert, + delete, + delete_, + detach, + + -- * Bisection methods + + -- ** C++-like + ilowerBound, + ilowerBoundM, + ilowerBoundProd, + ilowerBoundProdM, + + -- ** Splits + isplitMaxRight, + isplitMaxRightM, + isplitMaxRightProd, + isplitMaxRightProdM, + + -- * Conversions + freeze, + ) +where + +import AtCoder.Extra.Pool qualified as P +import AtCoder.Extra.Seq.Raw (Seq (..)) +import AtCoder.Extra.Seq.Raw qualified as Seq +import AtCoder.LazySegTree (SegAct (..)) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) +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, reverse, seq) + +-- | `Handle` for a sequence in `Seq`. It internally stores the root node and updates it +-- following splaying operations, as `Seq` utilizes a splay tree structure. +-- +-- @since 1.2.0.0 +newtype Handle s = Handle + { -- | @since 1.2.0.0 + unHandle :: VUM.MVector s P.Index + } + +-- | \(O(n)\) Creates a new `Seq` of length \(n\). +-- +-- @since 1.2.0.0 +{-# INLINE new #-} +new :: (PrimMonad m, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Int -> m (Seq (PrimState m) f a) +new n = stToPrim $ Seq.newST n + +-- | \(O(1)\) Creates a new sequence `Handle` from a root node index. +-- +-- @since 1.2.0.0 +{-# INLINE newHandle #-} +newHandle :: (PrimMonad m) => P.Index -> m (Handle (PrimState m)) +newHandle x = stToPrim $ Handle <$> VUM.replicate 1 x + +-- | \(O(1)\) Returns whether the sequence is empty. +-- +-- @since 1.2.0.0 +{-# INLINE nullHandle #-} +nullHandle :: (PrimMonad m) => Handle (PrimState m) -> m Bool +nullHandle (Handle h) = stToPrim $ do + P.nullIndex <$> VGM.unsafeRead h 0 + +-- | \(O(1)\) Invalidates a sequence handle. Note that it does not change or `free` the sequence. +-- +-- @since 1.2.0.0 +{-# INLINE invalidateHandle #-} +invalidateHandle :: (PrimMonad m) => Handle (PrimState m) -> m () +invalidateHandle (Handle h) = stToPrim $ do + VGM.unsafeWrite h 0 P.undefIndex + +-- | \(O(1)\) Clears the sequence storage. All the handles must not be used again. +-- +-- @since 1.2.0.0 +{-# INLINE reset #-} +reset :: (PrimMonad m) => Seq (PrimState m) f a -> m () +reset seq = stToPrim $ Seq.resetST seq + +-- | \(O(1)\) Allocates a new sequence of length \(1\). +-- +-- @since 1.2.0.0 +{-# INLINE newNode #-} +newNode :: (PrimMonad m, Monoid f, VU.Unbox f, VU.Unbox a) => Seq (PrimState m) f a -> a -> m (Handle (PrimState m)) +newNode seq x = stToPrim $ newHandle =<< Seq.newNodeST seq x + +-- | \(O(n)\) Allocates a new sequence. +-- +-- @since 1.2.0.0 +{-# INLINE newSeq #-} +newSeq :: (PrimMonad m, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> VU.Vector a -> m (Handle (PrimState m)) +newSeq seq !xs = stToPrim $ newHandle =<< Seq.newSeqST seq xs + +-- | \(O(n)\) Frees a sequence and invalidates the handle. +-- +-- @since 1.2.0.0 +{-# INLINE free #-} +free :: (PrimMonad m, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> m () +free seq (Handle handle) = stToPrim $ do + c0 <- VGM.unsafeRead handle 0 + Seq.freeSubtreeST seq c0 + VGM.write handle 0 P.undefIndex + +-- ------------------------------------------------------------------------------------------------- +-- Merge/split +-- ------------------------------------------------------------------------------------------------- + +-- | Amortized \(O(\log n)\). Merges two sequences \(l, r\) into one in the given order, ignoring +-- empty sequences. The right sequence handle will be invalidated. +-- +-- @since 1.2.0.0 +{-# INLINE merge #-} +merge :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Handle (PrimState m) -> m () +merge seq (Handle l) (Handle r) = stToPrim $ do + lRoot <- VGM.unsafeRead l 0 + rRoot <- VGM.unsafeRead r 0 + root' <- Seq.mergeST seq lRoot rRoot + VGM.unsafeWrite l 0 root' + VGM.unsafeWrite r 0 P.undefIndex + +-- | Amortized \(O(\log n)\). Merges three sequences \(l, m, r\) into one in the given order, +-- ignoring empty sequences. All handles, except for the leftmost one, will be invalidated. +-- +-- @since 1.2.0.0 +{-# INLINE merge3 #-} +merge3 :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Handle (PrimState m) -> Handle (PrimState m) -> m () +merge3 seq (Handle hA) (Handle hB) (Handle hC) = stToPrim $ do + a <- VGM.unsafeRead hA 0 + b <- VGM.unsafeRead hB 0 + c <- VGM.unsafeRead hC 0 + root' <- Seq.merge3ST seq a b c + VGM.unsafeWrite hA 0 root' + VGM.unsafeWrite hB 0 P.undefIndex + VGM.unsafeWrite hC 0 P.undefIndex + +-- | Amortized \(O(\log n)\). Merges four sequences \(a, b, c, d\) into one in the given order, +-- ignoring empty sequences. All handles, except for the leftmost one, will be invalidated. +-- +-- ==== Constraints +-- - The vertices must be roots. +-- +-- @since 1.2.0.0 +{-# INLINE merge4 #-} +merge4 :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Handle (PrimState m) -> Handle (PrimState m) -> Handle (PrimState m) -> m () +merge4 seq (Handle hA) (Handle hB) (Handle hC) (Handle hD) = stToPrim $ do + a <- VGM.unsafeRead hA 0 + b <- VGM.unsafeRead hB 0 + c <- VGM.unsafeRead hC 0 + d <- VGM.unsafeRead hC 0 + root' <- Seq.merge4ST seq a b c d + VGM.unsafeWrite hA 0 root' + VGM.unsafeWrite hB 0 P.undefIndex + VGM.unsafeWrite hC 0 P.undefIndex + VGM.unsafeWrite hD 0 P.undefIndex + +-- | Amortized \(O(\log n)\). Splits a sequences into two: \([0, k), [k, n)\). The handle will +-- point to the left sequence. Returns the right sequence handle. +-- +-- ==== Constraints +-- - \(0 \le k \le n\). +-- +-- @since 1.2.0.0 +{-# INLINE split #-} +split :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> m (Handle (PrimState m)) +split seq (Handle hRoot) k = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + (!r1, !r2) <- Seq.splitST seq root k + VGM.unsafeWrite hRoot 0 r1 + newHandle r2 + +-- | Amortized \(O(\log n)\). Splits a sequences into three: \([0, l), [l, r), [r, n)\). The handle +-- will point to the leftmost sequence. Returns the middle and the right sequence handles. +-- +-- ==== Constraints +-- - \(0 \le l \le r \le n\). +-- +-- @since 1.2.0.0 +{-# INLINE split3 #-} +split3 :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> Int -> m (Handle (PrimState m), Handle (PrimState m)) +split3 seq (Handle hRoot) l r = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + (!r1, !r2, !r3) <- Seq.split3ST seq root l r + VGM.unsafeWrite hRoot 0 r1 + (,) <$> newHandle r2 <*> newHandle r3 + +-- | Amortized \(O(\log n)\). Splits a sequences into four: \([0, i), [i, j), [j, k), [k, n)\). +-- The handle will point to the leftmost sequence. Returns the non-leftmost sequence handles. +-- +-- ==== Constraints +-- - \(0 \le i \le j \le k \le n\). +-- +-- @since 1.2.0.0 +{-# INLINE split4 #-} +split4 :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> Int -> Int -> m (Handle (PrimState m), Handle (PrimState m), Handle (PrimState m)) +split4 seq (Handle hRoot) i j k = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + (!r1, !r2, !r3, !r4) <- Seq.split4ST seq root i j k + VGM.unsafeWrite hRoot 0 r1 + (,,) <$> newHandle r2 <*> newHandle r3 <*> newHandle r4 + +-- | Amortized \(O(\log n)\). Splits a sequence into three: \([0, \mathrm{root}), \mathrm{root}, [\mathrm{root} + 1, n)\). +-- +-- ==== Constraints +-- - The sequence must be non-empty. +-- +-- @since 1.2.0.0 +{-# INLINE splitLr #-} +splitLr :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> m (Handle (PrimState m), Handle (PrimState m)) +splitLr seq (Handle hRoot) = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + (!l, !root', !r) <- Seq.splitLrST seq root + VGM.unsafeWrite hRoot 0 root' + (,) <$> newHandle l <*> newHandle r + +-- ------------------------------------------------------------------------------------------------- +-- Modifications +-- ------------------------------------------------------------------------------------------------- + +-- | Amortized \(O(\log n)\). Reads the \(k\)-th node's monoid value. +-- +-- ==== Constraints +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINE read #-} +read :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> m a +read seq (Handle hRoot) k = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + (!v, !root') <- Seq.readST seq root k + VGM.unsafeWrite hRoot 0 root' + pure v + +-- | Amortized \(O(\log n)\). Reads the \(k\)-th node's monoid value. +-- +-- @since 1.2.0.0 +{-# INLINE readMaybe #-} +readMaybe :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> m (Maybe a) +readMaybe seq (Handle hRoot) k = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + res <- Seq.readMaybeST seq root k + case res of + Just (!v, !root') -> do + VGM.unsafeWrite hRoot 0 root' + pure $ Just v + Nothing -> pure Nothing + +-- | Amortized \(O(\log n)\). Writes to the \(k\)-th node's monoid value. +-- +-- ==== Constraints +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINE write #-} +write :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> a -> m () +write seq (Handle hRoot) k v = stToPrim $ do + VGM.unsafeModifyM + hRoot + ( \root -> do + Seq.writeST seq root k v + ) + 0 + +-- | Amortized \(O(\log n)\). Modifies the \(k\)-th node's monoid value. +-- +-- ==== Constraints +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINE modify #-} +modify :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> (a -> a) -> Int -> m () +modify seq (Handle hRoot) f k = stToPrim $ do + VGM.unsafeModifyM + hRoot + ( \root -> do + Seq.modifyST seq root f k + ) + 0 + +-- | Amortized \(O(\log n)\). Exchanges the \(k\)-th node's monoid value. +-- +-- ==== Constraints +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINE exchange #-} +exchange :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> a -> m a +exchange seq (Handle hRoot) k v = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + (!x, !root') <- Seq.exchangeST seq root k v + VGM.unsafeWrite hRoot 0 root' + pure x + +-- | Amortized \(O(\log n)\). Returns the monoid product in an interval \([l, r)\). +-- +-- ==== Constraints +-- - \(0 \le l \le r \le n\) +-- +-- @since 1.2.0.0 +{-# INLINE prod #-} +prod :: (HasCallStack, Show a, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> Int -> m a +prod seq (Handle hRoot) l r = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + (!v, !root') <- Seq.prodST seq root l r + VGM.unsafeWrite hRoot 0 root' + pure v + +-- | Amortized \(O(\log n)\). Returns the monoid product in an interval \([l, r)\). Returns +-- `Nothing` if an invalid interval is given. +-- +-- @since 1.2.0.0 +{-# INLINE prodMaybe #-} +prodMaybe :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> Int -> m (Maybe a) +prodMaybe seq (Handle handle) l r = stToPrim $ do + root <- VGM.unsafeRead handle 0 + res <- Seq.prodMaybeST seq root l r + case res of + Just (!v, !root') -> do + VGM.unsafeWrite handle 0 root' + pure $ Just v + Nothing -> pure Nothing + +-- | Amortized \(O(\log n)\). Returns the monoid product of the whole sequence. +-- +-- @since 1.2.0.0 +{-# INLINE prodAll #-} +prodAll :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> m a +prodAll seq (Handle handle) = stToPrim $ do + root <- VGM.unsafeRead handle 0 + Seq.prodAllST seq root + +-- | Amortized \(O(\log n)\). Given an interval \([l, r)\), applies a monoid action \(f\). +-- +-- ==== Constraints +-- - \(0 \le l \le r \le n\) +-- +-- @since 1.2.0.0 +{-# INLINE applyIn #-} +applyIn :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> Int -> f -> m () +applyIn seq (Handle hRoot) l r act = stToPrim $ do + VGM.unsafeModifyM + hRoot + ( \root -> do + Seq.applyInST seq root l r act + ) + 0 + +-- | \(O(1)\) Applies a monoid action \(f\) to the root of a sequence. +-- +-- @since 1.2.0.0 +{-# INLINE applyToRoot #-} +applyToRoot :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> f -> m () +applyToRoot seq (Handle hRoot) act = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + Seq.applyToRootST seq root act + +-- | Amortized \(O(\log n)\). Reverses the sequence in \([l, r)\). +-- +-- ==== Constraints +-- - The monoid action \(f\) must be commutative. +-- - The monoid value \(v\) must be commutative. +-- +-- @since 1.2.0.0 +{-# INLINE reverse #-} +reverse :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> Int -> m () +reverse seq (Handle hRoot) l r = stToPrim $ do + VGM.unsafeModifyM + hRoot + ( \root -> do + Seq.reverseST seq root l r + ) + 0 + +-- | Amortized \(O(\log n)\). Inserts a new node at \(k\) with initial monoid value \(v\). This +-- function works for an empty sequence handle. +-- +-- ==== Constraints +-- - \(0 \le k \le n\) +-- +-- @since 1.2.0.0 +{-# INLINE insert #-} +insert :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> a -> m () +insert seq (Handle hRoot) k v = stToPrim $ do + VGM.unsafeModifyM + hRoot + ( \root -> do + Seq.insertST seq root k v + ) + 0 + +-- | Amortized \(O(\log n)\). Frees the \(k\)-th node and returns the monoid value of it. +-- +-- ==== Constraints +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINE delete #-} +delete :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> m a +delete seq (Handle hRoot) i = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + (!v, !root') <- Seq.deleteST seq root i + VGM.unsafeWrite hRoot 0 root' + pure v + +-- | Amortized \(O(\log n)\). Frees the \(k\)-th node. +-- +-- ==== Constraints +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINE delete_ #-} +delete_ :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> m () +delete_ seq (Handle hRoot) i = stToPrim $ do + VGM.unsafeModifyM + hRoot + ( \root -> do + Seq.deleteST_ seq root i + ) + 0 + +-- | Amortized \(O(\log n)\). Detaches the \(k\)-th node and returns a handle for it. +-- +-- ==== Constraints +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINE detach #-} +detach :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> Int -> m (Handle (PrimState m)) +detach seq (Handle hRoot) i = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + root' <- Seq.detachST seq root i + VGM.unsafeWrite hRoot 0 root' + newHandle root + +-- ------------------------------------------------------------------------------------------------- +-- Bisection methods +-- ------------------------------------------------------------------------------------------------- + +-- | Amortized \(O(\log n)\). +-- +-- ==== Constraints +-- - The sequence must be non-empty. +-- +-- @since 1.2.0.0 +{-# INLINE ilowerBound #-} +ilowerBound :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Sequence handle + Handle (PrimState m) -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> Bool) -> + -- | Maximum \(r\), where \(f(i, v_i)\) holds for \(i \in [0, r)\) + m Int +ilowerBound seq (Handle root0) f = stToPrim $ do + root <- VGM.unsafeRead root0 0 + (!r, !root') <- Seq.ilowerBoundST seq root f + VGM.unsafeWrite root0 0 root' + pure r + +-- | Amortized \(O(\log n)\). +-- +-- ==== Constraints +-- - The sequence must be non-empty. +-- +-- @since 1.2.0.0 +{-# INLINE ilowerBoundM #-} +ilowerBoundM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Sequence handle + Handle (PrimState m) -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> m Bool) -> + -- | Maximum \(r\), where \(f(i, v_i)\) holds for \(i \in [0, r)\) + m Int +ilowerBoundM seq (Handle root0) f = do + root <- VGM.unsafeRead root0 0 + (!r, !root') <- Seq.ilowerBoundM seq root f + VGM.unsafeWrite root0 0 root' + pure r + +-- | Amortized \(O(\log n)\). +-- +-- ==== Constraints +-- - The sequence must be non-empty. +-- +-- @since 1.2.0.0 +{-# INLINE ilowerBoundProd #-} +ilowerBoundProd :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Sequence handle + Handle (PrimState m) -> + -- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid product + (Int -> a -> Bool) -> + -- | Maximum \(r\), where \(f(i, v_0 \dots v_i)\) holds for \(i \in [0, r)\) + m Int +ilowerBoundProd seq (Handle root0) f = stToPrim $ do + root <- VGM.unsafeRead root0 0 + (!r, !root') <- Seq.ilowerBoundProdST seq root f + VGM.unsafeWrite root0 0 root' + pure r + +-- | Amortized \(O(\log n)\). +-- +-- ==== Constraints +-- - The sequence must be non-empty. +-- +-- @since 1.2.0.0 +{-# INLINE ilowerBoundProdM #-} +ilowerBoundProdM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Sequence handle + Handle (PrimState m) -> + -- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid product + (Int -> a -> m Bool) -> + -- | Maximum \(r\), where \(f(i, v_0 \dots v_i)\) holds for \(i \in [0, r)\) + m Int +ilowerBoundProdM seq (Handle root0) f = do + root <- VGM.unsafeRead root0 0 + (!r, !root') <- Seq.ilowerBoundProdM seq root f + VGM.unsafeWrite root0 0 root' + pure r + +-- | Amortized \(O(\log n)\). Splits a sequence into two with the user predicate and returns the +-- right sequence handle. +-- +-- ==== Constraints +-- - The sequence must be non-empty. +-- +-- @since 1.2.0.0 +{-# INLINE isplitMaxRight #-} +isplitMaxRight :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Sequence handle + Handle (PrimState m) -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> Bool) -> + -- | Handle of the right sequence \([r, n)\), where \(r\) is the maximum \(r\) such that + -- \(f(i, v_i)\) holds for \(i \in [0, r)\) + m (Handle (PrimState m)) +isplitMaxRight seq (Handle root0) f = stToPrim $ do + root <- VGM.unsafeRead root0 0 + (!l, !r) <- Seq.isplitMaxRightST seq root f + VGM.unsafeWrite root0 0 l + newHandle r + +-- | Amortized \(O(\log n)\). Splits a sequence into two with the user predicate and returns the +-- right sequence handle. +-- +-- ==== Constraints +-- - The sequence must be non-empty. +-- +-- @since 1.2.0.0 +{-# INLINEABLE isplitMaxRightM #-} +isplitMaxRightM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Sequence handle + Handle (PrimState m) -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> m Bool) -> + -- | Handle of the right sequence \([r, n)\), where \(r\) is the maximum \(r\) such that + -- \(f(i, v_i)\) holds for \(i \in [0, r)\) + m (Handle (PrimState m)) +isplitMaxRightM seq (Handle root0) f = do + root <- VGM.unsafeRead root0 0 + (!l, !r) <- Seq.isplitMaxRightM seq root f + VGM.unsafeWrite root0 0 l + newHandle r + +-- | Amortized \(O(\log n)\). Splits a sequence into two with the user predicate and returns the +-- right sequence handle. +-- +-- ==== Constraints +-- - The sequence must be non-empty. +-- +-- @since 1.2.0.0 +{-# INLINE isplitMaxRightProd #-} +isplitMaxRightProd :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Sequence handle + Handle (PrimState m) -> + -- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid value + (Int -> a -> Bool) -> + -- | Handle of the right sequence \([r, n)\), where \(r\) is the maximum \(r\) such that + -- \(f(i, v_0 \dots v_i)\) holds for \(i \in [0, r)\) + m (Handle (PrimState m)) +isplitMaxRightProd seq (Handle root0) f = stToPrim $ do + root <- VGM.unsafeRead root0 0 + (!l, !r) <- Seq.isplitMaxRightProdST seq root f + VGM.unsafeWrite root0 0 l + newHandle r + +-- | Amortized \(O(\log n)\). Splits a sequence into two with the user predicate and returns the +-- right sequence handle. +-- +-- ==== Constraints +-- - The sequence must be non-empty. +-- +-- @since 1.2.0.0 +{-# INLINEABLE isplitMaxRightProdM #-} +isplitMaxRightProdM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Sequence handle + Handle (PrimState m) -> + -- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid value + (Int -> a -> m Bool) -> + -- | Handle of the right sequence \([r, n)\), where \(r\) is the maximum \(r\) such that + -- \(f(i, v_0 \dots v_i)\) holds for \(i \in [0, r)\) + m (Handle (PrimState m)) +isplitMaxRightProdM seq (Handle root0) f = do + root <- VGM.unsafeRead root0 0 + (!l, !r) <- Seq.isplitMaxRightProdM seq root f + VGM.unsafeWrite root0 0 l + newHandle r + +-- ------------------------------------------------------------------------------------------------- +-- Conversions +-- ------------------------------------------------------------------------------------------------- + +-- | Amortized \(O(n)\). Returns the sequence of monoid values. +-- +-- @since 1.2.0.0 +{-# INLINEABLE freeze #-} +freeze :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq (PrimState m) f a -> Handle (PrimState m) -> m (VU.Vector a) +freeze seq (Handle hRoot) = stToPrim $ do + root <- VGM.unsafeRead hRoot 0 + Seq.freezeST seq root diff --git a/src/AtCoder/Extra/Seq/Raw.hs b/src/AtCoder/Extra/Seq/Raw.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/Seq/Raw.hs @@ -0,0 +1,1295 @@ +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Base module for implementing dynamic sequences. It internaly uses a splay tree and user has to +-- track the root node change. +-- +-- @since 1.2.0.0 +module AtCoder.Extra.Seq.Raw + ( -- * Seq + Seq (..), + + -- * Constructors + newST, + resetST, + newNodeST, + newSeqST, + freeNodeST, + freeSubtreeST, + + -- * Merge/split + mergeST, + merge3ST, + merge4ST, + splitST, + split3ST, + split4ST, + splitLrST, + sliceST, + + -- * Read/write + readST, + readMaybeST, + writeST, + modifyST, + exchangeST, + + -- * Products + prodST, + prodMaybeST, + prodAllST, + + -- * Applications + applyInST, + applyToRootST, + reverseST, + + -- * Insert/delete + insertST, + deleteST, + deleteST_, + detachST, + + -- * Balancing + rotateST, + splayST, + splayKthST, + + -- * Bisection methods + + -- ** C++-like + ilowerBoundST, + ilowerBoundM, + ilowerBoundProdST, + ilowerBoundProdM, + + -- ** Splits + isplitMaxRightST, + isplitMaxRightM, + isplitMaxRightProdST, + isplitMaxRightProdM, + + -- ** Max right + imaxRightST, + imaxRightM, + imaxRightProdST, + imaxRightProdM, + + -- * Conversions + freezeST, + ) +where + +import AtCoder.Extra.Pool qualified as P +import AtCoder.Internal.Assert qualified as ACIA +import AtCoder.LazySegTree (SegAct (..)) +import Control.Monad (unless, when) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) +import Control.Monad.ST (ST) +import Data.Bit +import Data.Bits hiding (rotate) +import Data.Coerce (coerce) +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 (seq) + +-- | Storages of dynamic sequences of monoid values with monoid actions on them through the `SegAct` +-- instance. +-- +-- @since 1.2.0.0 +data Seq s f a = Seq + { -- | The maximum number of elements. + -- + -- @since 1.2.0.0 + nSeq :: {-# UNPACK #-} !Int, + -- | `Pool` for free slot management. + -- + -- @since 1.2.0.0 + poolSeq :: !(P.Pool s ()), + -- | Decomposed node data storage: left children. + -- + -- @since 1.2.0.0 + lSeq :: !(VUM.MVector s P.Index), + -- | Decomposed node data storage: right children. + -- + -- @since 1.2.0.0 + rSeq :: !(VUM.MVector s P.Index), + -- | Decomposed node data storage: parents. + -- + -- @since 1.2.0.0 + pSeq :: !(VUM.MVector s P.Index), + -- | Decomposed node data storage: subtree sizes. + -- + -- @since 1.2.0.0 + sSeq :: !(VUM.MVector s Int), + -- | Decomposed node data storage: monoid values. + -- + -- @since 1.2.0.0 + vSeq :: !(VUM.MVector s a), + -- | Decomposed node data storage: monoid products. + -- + -- @since 1.2.0.0 + prodSeq :: !(VUM.MVector s a), + -- | Decomposed node data storage: reversed flag of children. + -- + -- @since 1.2.0.0 + revSeq :: !(VUM.MVector s Bit), + -- | Decomposed node data storage: lazily propagated monoid action. Use @()@ if you don't need + -- monoid actions. + -- + -- @since 1.2.0.0 + lazySeq :: !(VUM.MVector s f) + } + +-- | \(O(n)\) Creates a new `Seq` of length \(n\). +-- +-- @since 1.2.0.0 +{-# INLINEABLE newST #-} +newST :: (Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Int -> ST s (Seq s f a) +newST nSeq = do + poolSeq <- P.new nSeq + lSeq <- VUM.unsafeNew nSeq + rSeq <- VUM.unsafeNew nSeq + pSeq <- VUM.unsafeNew nSeq + sSeq <- VUM.unsafeNew nSeq + vSeq <- VUM.unsafeNew nSeq + prodSeq <- VUM.unsafeNew nSeq + revSeq <- VUM.unsafeNew nSeq + lazySeq <- VUM.unsafeNew nSeq + pure Seq {..} + +-- | \(O(1)\) Clears the sequence storage. +-- +-- @since 1.2.0.0 +{-# INLINE resetST #-} +resetST :: Seq s f a -> ST s () +resetST Seq {poolSeq} = stToPrim $ P.clear poolSeq + +-- | \(O(1)\) Allocates a new sequence of length \(1\). +-- +-- @since 1.2.0.0 +{-# INLINEABLE newNodeST #-} +newNodeST :: (Monoid f, VU.Unbox f, VU.Unbox a) => Seq s f a -> a -> ST s P.Index +newNodeST Seq {..} x = do + i <- P.alloc poolSeq () + VGM.write lSeq (coerce i) P.undefIndex + VGM.write rSeq (coerce i) P.undefIndex + VGM.write pSeq (coerce i) P.undefIndex + VGM.write sSeq (coerce i) 1 + VGM.write vSeq (coerce i) x + VGM.write prodSeq (coerce i) x + VGM.write revSeq (coerce i) $ Bit False + VGM.write lazySeq (coerce i) mempty + pure i + +-- | \(O(n)\) Allocates a new sequence. +-- +-- @since 1.2.0.0 +{-# INLINEABLE newSeqST #-} +newSeqST :: (Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> VU.Vector a -> ST s P.Index +newSeqST seq@Seq {..} !xs = do + -- [l, r) + let inner l r + | l >= r = pure P.undefIndex + | l + 1 == r = newNodeST seq $ xs VG.! l + | otherwise = do + let !m = (l + r) `div` 2 + rootL <- inner l m + rootR <- inner (m + 1) r + root <- newNodeST seq (xs VG.! m) + unless (P.nullIndex rootL) $ do + VGM.write lSeq (coerce root) rootL + VGM.write pSeq (coerce rootL) root + unless (P.nullIndex rootR) $ do + VGM.write rSeq (coerce root) rootR + VGM.write pSeq (coerce rootR) root + updateNodeST seq root + pure root + inner 0 (VU.length xs) + +-- | \(O(1)\) Frees a node. +-- +-- @since 1.2.0.0 +{-# INLINE freeNodeST #-} +freeNodeST :: Seq s v a -> P.Index -> ST s () +freeNodeST Seq {poolSeq} = P.free poolSeq + +-- | \(O(n)\) Frees a subtree. +-- +-- @since 1.2.0.0 +{-# INLINEABLE freeSubtreeST #-} +freeSubtreeST :: (VU.Unbox a) => Seq s f a -> P.Index -> ST s () +freeSubtreeST Seq {lSeq, rSeq, poolSeq} c0 + | P.nullIndex c0 = pure () + | otherwise = do + let inner c = do + cl <- VGM.read lSeq (coerce c) + unless (P.nullIndex cl) (inner cl) + cr <- VGM.read rSeq (coerce c) + unless (P.nullIndex cr) (inner cr) + inner c0 + P.free poolSeq c0 + +-- ------------------------------------------------------------------------------------------------- +-- Merge/split +-- ------------------------------------------------------------------------------------------------- + +{-# INLINE assertRootST #-} +assertRootST :: (HasCallStack) => Seq s f a -> P.Index -> ST s () +assertRootST Seq {pSeq} i = do + p <- VGM.read pSeq (coerce i) + let !_ = ACIA.runtimeAssert (P.nullIndex p) $ "AtCoder.Extra.Seq.Raw.assertRootST: not a root (node `" ++ show i ++ "`, parent `" ++ show p ++ "`)" + pure () + +{-# INLINE assertRootOrNullST #-} +assertRootOrNullST :: (HasCallStack) => Seq s f a -> P.Index -> ST s () +assertRootOrNullST Seq {pSeq} i + | P.nullIndex i = pure () + | otherwise = do + p <- VGM.read pSeq (coerce i) + let !_ = ACIA.runtimeAssert (P.nullIndex p) $ "AtCoder.Extra.Seq.Raw.assertRootOrNullST: not a root (node `" ++ show i ++ "`, parent `" ++ show p ++ "`)" + pure () + +-- | Amortized \(O(\log n)\). Merges two sequences \(l, r\) into one in the given order, ignoring +-- empty sequences. +-- +-- ==== Constraints +-- - The vertices must be either null or a root. +-- +-- @since 1.2.0.0 +{-# INLINEABLE mergeST #-} +mergeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> P.Index -> ST s P.Index +mergeST seq@Seq {pSeq, lSeq} lRoot rRoot + | P.nullIndex lRoot = pure rRoot + | P.nullIndex rRoot = pure lRoot + | otherwise = do + do + -- TODO: delete + lp <- VGM.read pSeq (coerce lRoot) + rp <- VGM.read pSeq (coerce rRoot) + let !_ = ACIA.runtimeAssert (lp == rp) "AtCoder.Extra.Seq.Raw.mergeST: given non-root node" + pure () + rRoot' <- splayKthST seq rRoot 0 + VGM.write lSeq (coerce rRoot') lRoot + VGM.write pSeq (coerce lRoot) rRoot' + updateNodeST seq rRoot' + pure rRoot' + +-- | Amortized \(O(\log n)\). Merges three sequences \(l, m, r\) into one in the given order, +-- ignoring empty sequences. +-- +-- ==== Constraints +-- - The vertices must be either null or a root. +-- +-- @since 1.2.0.0 +{-# INLINE merge3ST #-} +merge3ST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> P.Index -> P.Index -> ST s P.Index +merge3ST seq a b c = do + r' <- mergeST seq a b + mergeST seq r' c + +-- | Amortized \(O(\log n)\). Merges four sequences \(l, b, c, d, m, r\) into one in the given +-- order, ignoring empty sequences. +-- +-- ==== Constraints +-- - The vertices must be either null or a root. +-- +-- @since 1.2.0.0 +{-# INLINE merge4ST #-} +merge4ST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> P.Index -> P.Index -> P.Index -> ST s P.Index +merge4ST seq a b c d = do + r' <- mergeST seq a b + r'' <- mergeST seq r' c + mergeST seq r'' d + +-- | Amortized \(O(\log n)\). Splits a sequences into two: \([0, k), [k, n)\). +-- +-- ==== Constraints +-- - The node must be null or a root. +-- - \(0 \le k \le n\). +-- +-- @since 1.2.0.0 +{-# INLINEABLE splitST #-} +splitST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s (P.Index, P.Index) +splitST seq@Seq {..} root k = do + assertRootOrNullST seq root + if k == 0 + then pure (P.undefIndex, root) + else do + size <- VGM.read sSeq $ coerce root + if k == size + then pure (root, P.undefIndex) + else do + root' <- splayKthST seq root (k - 1) + r <- VGM.exchange rSeq (coerce root') P.undefIndex + VGM.write pSeq (coerce r) P.undefIndex + updateNodeST seq root' + pure (root', r) + +-- | Amortized \(O(\log n)\). Splits a sequences into three: \([0, l), [l, r), [r, n)\). +-- +-- ==== Constraints +-- - The node must be null or a root. +-- - \(0 \le l \le r \le n\). +-- +-- @since 1.2.0.0 +{-# INLINE split3ST #-} +split3ST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s (P.Index, P.Index, P.Index) +split3ST seq root l r = do + (!root', !c) <- splitST seq root r + (!a, !b) <- splitST seq root' l + pure (a, b, c) + +-- | Amortized \(O(\log n)\). Splits a sequences into four: \([0, i), [i, j), [j, k), [k, n)\). +-- +-- ==== Constraints +-- - The node must be null or a root. +-- - \(0 \le i \le j \le k \le n\). +-- +-- @since 1.2.0.0 +{-# INLINE split4ST #-} +split4ST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> Int -> ST s (P.Index, P.Index, P.Index, P.Index) +split4ST seq root i j k = do + (!root', !d) <- splitST seq root k + (!root'', !c) <- splitST seq root' j + (!a, !b) <- splitST seq root'' i + pure (a, b, c, d) + +-- | Amortized \(O(\log n)\). Splits a sequence into three: \([0, \mathrm{root}), \mathrm{root}, [\mathrm{root} + 1, n)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINEABLE splitLrST #-} +splitLrST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s (P.Index, P.Index, P.Index) +splitLrST seq@Seq {..} root = do + assertRootST seq root + s <- do + rootL <- VGM.read lSeq (coerce root) + if P.nullIndex rootL + then VGM.read sSeq (coerce rootL) + else pure 0 + split3ST seq root s (s + 1) + +-- | Amortized \(O(\log n)\). Captures the root of a subtree of \([l, r)\). Splay the new root after +-- call. +-- +-- ==== Constraints +-- - \(0 \le \lt r \le n\). Note that the interval must have positive length. +-- +-- @since 1.2.0.0 +{-# INLINEABLE sliceST #-} +sliceST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s P.Index +sliceST seq@Seq {..} root l r + | l == 0 = do + size <- VGM.read sSeq (coerce root) + if r == size + then pure root + else do + root' <- splayKthST seq root r + VGM.read lSeq $ coerce root' + | otherwise = do + size <- VGM.read sSeq $ coerce root + if r == size + then do + root' <- splayKthST seq root (l - 1) + VGM.read rSeq $ coerce root' + else do + -- o--l--o--o--r--o + -- [ ) + -- * root' (splayed) + -- * rootL (detached from the root) + -- \* rootL' (splayed) + -- * right(rootL'): node that corresponds to [l, r) + root' <- splayKthST seq root r + rootL <- VGM.read lSeq $ coerce root' + -- detach `rootL` from `root'` + VGM.write pSeq (coerce rootL) P.undefIndex + rootL' <- splayKthST seq rootL (l - 1) + -- re-attach `rootL'` to `root'` + VGM.write pSeq (coerce rootL') root' + VGM.write lSeq (coerce root') rootL' + updateNodeST seq root' + VGM.read rSeq $ coerce rootL' + +-- ------------------------------------------------------------------------------------------------- +-- Modifications +-- ------------------------------------------------------------------------------------------------- + +-- | Amortized \(O(\log n)\). Reads the \(k\)-th node's monoid value. +-- +-- ==== Constraints +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE readST #-} +readST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s (a, P.Index) +readST seq@Seq {..} root k = do + assertRootST seq root + root' <- splayKthST seq root k + (,root') <$> VGM.read vSeq (coerce root') + +-- | Amortized \(O(\log n)\). Reads the \(k\)-th node's monoid value. +-- +-- ==== Constraints +-- - The root must be empty or a root. +-- +-- @since 1.2.0.0 +{-# INLINEABLE readMaybeST #-} +readMaybeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s (Maybe (a, P.Index)) +readMaybeST seq@Seq {..} root k + | P.nullIndex root = pure Nothing + | otherwise = do + assertRootST seq root + s <- VGM.read sSeq (coerce root) + if 0 <= k && k < s + then do + root' <- splayKthST seq root k + Just . (,root') <$> VGM.read vSeq (coerce root') + else pure Nothing + +-- | Amortized \(O(\log n)\). Writes to the \(k\)-th node's monoid value. +-- +-- ==== Constraints +-- - The node must be a root. +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE writeST #-} +writeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> a -> ST s P.Index +writeST seq root k v = do + assertRootST seq root + root' <- splayKthST seq root k + writeNodeST seq root' v + pure root' + +-- | Amortized \(O(\log n)\). Modifies the \(k\)-th node's monoid value. +-- +-- ==== Constraints +-- - The node must be a root. +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE modifyST #-} +modifyST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> (a -> a) -> Int -> ST s P.Index +modifyST seq root f k = do + assertRootST seq root + root' <- splayKthST seq root k + modifyNodeST seq f root' + pure root' + +-- | Amortized \(O(\log n)\). Exchanges the \(k\)-th node's monoid value. +-- +-- ==== Constraints +-- - The node must be a root. +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE exchangeST #-} +exchangeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> a -> ST s (a, P.Index) +exchangeST seq root k v = do + assertRootST seq root + root' <- splayKthST seq root k + res <- exchangeNodeST seq root' v + pure (res, root') + +-- | Amortized \(O(\log n)\). Returns the monoid product in an interval \([l, r)\). +-- +-- ==== Constraints +-- - The node must be a root +-- - \(0 \le l \le r \le n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE prodST #-} +prodST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s (a, P.Index) +prodST seq@Seq {sSeq} root l r = do + s <- if P.nullIndex root then pure 0 else VGM.read sSeq (coerce root) + let !_ = ACIA.checkInterval "AtCoder.Extra.Seq.Raw.prodST" l r s + if l == r + then pure (mempty, root) + else unsafeProdST seq root l r + +-- | Amortized \(O(\log n)\). Returns the monoid product in an interval \([l, r)\). Returns +-- `Nothing` if an invalid interval is given or for an empty sequence. +-- +-- @since 1.2.0.0 +{-# INLINEABLE prodMaybeST #-} +prodMaybeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s (Maybe (a, P.Index)) +prodMaybeST seq@Seq {sSeq} root l r + | P.nullIndex root = pure Nothing + | otherwise = do + s <- VGM.read sSeq (coerce root) + if not (ACIA.testInterval l r s) + then pure Nothing + else + if l == r + then pure $ Just (mempty, root) + else Just <$> unsafeProdST seq root l r + +-- | Amortized \(O(\log n)\). +-- +-- ==== Constraint +-- - \(0 \le \lt r \le n\). Note that the interval must have positive length. +{-# INLINEABLE unsafeProdST #-} +unsafeProdST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s (a, P.Index) +unsafeProdST seq@Seq {..} root l r = do + assertRootST seq root + target <- sliceST seq root l r + res <- VGM.read prodSeq $ coerce target + splayST seq target True + pure (res, target) + +-- | Amortized \(O(\log n)\). Returns the monoid product of the whole sequence. Returns `mempty` +-- for an empty sequence. +-- +-- ==== Constraint +-- - The node must be null or a root. +-- +-- @since 1.2.0.0 +{-# INLINEABLE prodAllST #-} +prodAllST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s a +prodAllST seq@Seq {..} root = do + if P.nullIndex root + then pure mempty + else do + assertRootST seq root + VGM.read prodSeq $ coerce root + +-- | Amortized \(O(\log n)\). Given an interval \([l, r)\), applies a monoid action \(f\). +-- +-- ==== Constraints +-- - \(0 \le l \le r \le n\) +-- - The root must point to a non-empty sequence. +-- +-- @since 1.2.0.0 +{-# INLINEABLE applyInST #-} +applyInST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> f -> ST s P.Index +applyInST seq@Seq {..} root l r act = do + assertRootST seq root + s <- if P.nullIndex root then pure 0 else VGM.read sSeq (coerce root) + let !_ = ACIA.checkInterval "AtCoder.Extra.Seq.applyInST" l r s + if l == r + then pure root + else do + root' <- sliceST seq root l r + applyNodeST seq root' act + splayST seq root' True + pure root' + +-- | \(O(1)\) Applies a monoid action \(f\) to the root of a sequence. +-- +-- @since 1.2.0.0 +{-# INLINEABLE applyToRootST #-} +applyToRootST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> f -> ST s () +applyToRootST seq@Seq {..} root act + | P.nullIndex root = pure () + | otherwise = do + rootP <- VGM.read pSeq (coerce root) + when (P.nullIndex rootP) $ do + applyNodeST seq root act + +-- | Amortized \(O(\log n)\). Reverses the sequence in \([l, r)\). +-- +-- ==== Constraints +-- - The monoid action \(f\) must be commutative. +-- - The monoid value \(v\) must be commutative. +-- +-- @since 1.2.0.0 +{-# INLINEABLE reverseST #-} +reverseST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s P.Index +reverseST seq@Seq {sSeq} root0 l r + | P.nullIndex root0 = pure P.undefIndex + | otherwise = do + s <- VGM.read sSeq (coerce root0) + if not (ACIA.testInterval l r s) + then pure root0 + else + if l == r + then pure root0 + else do + root' <- sliceST seq root0 l r + reverseNodeST seq root' + splayST seq root' True + pure root' + +-- | Amortized \(O(\log n)\). Inserts a new node at \(k\) with initial monoid value \(v\). This +-- functions for an empty index. +-- +-- ==== Constraints +-- - The node must be null or a root. +-- - \(0 \le k \le n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE insertST #-} +insertST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> a -> ST s P.Index +insertST seq root k v = do + if P.nullIndex root + then do + -- `insertST` is actually `insertOrNewNodeST`: it's specifically designed to work for an empty + -- sequence. + newNodeST seq v + else do + (!l, !r) <- splitST seq root k + node <- newNodeST seq v + merge3ST seq l node r + +-- | Amortized \(O(\log n)\). Frees the \(k\)-th node and returns the monoid value of it. +-- +-- ==== Constraints +-- - The node must be null or a root. +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE deleteST #-} +deleteST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s (a, P.Index) +deleteST seq@Seq {..} root i = do + (!l, !m, !r) <- split3ST seq root i (i + 1) + x <- VGM.read vSeq (coerce m) + freeNodeST seq m + root' <- mergeST seq l r + pure (x, root') + +-- | Amortized \(O(\log n)\). Frees the \(k\)-th node. +-- +-- ==== Constraints +-- - The node must be null or a root. +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE deleteST_ #-} +deleteST_ :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s P.Index +deleteST_ seq root i = do + (!l, !m, !r) <- split3ST seq root i (i + 1) + freeNodeST seq m + root' <- mergeST seq l r + pure root' + +-- | Amortized \(O(\log n)\). Detaches the \(k\)-th node and returns the new root of the original +-- sequence. +-- +-- ==== Constraints +-- - The node must be null or a root. +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE detachST #-} +detachST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s P.Index +detachST seq root i = do + (!l, !m, !r) <- split3ST seq root i (i + 1) + freeNodeST seq m + root' <- mergeST seq l r + pure root' + +-- ------------------------------------------------------------------------------------------------- +-- Balancing +-- ------------------------------------------------------------------------------------------------- + +-- | Amortized \(O(\log n)\). Rotates a child node. +-- +-- ==== Constraints +-- - \(0 \le i \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE rotateST #-} +rotateST :: (HasCallStack) => Seq s v a -> P.Index -> ST s () +rotateST Seq {..} !i = do + p <- VGM.read pSeq $ coerce i + pl <- VGM.read lSeq $ coerce p + + c <- + if pl == i + then do + -- p i + -- / \ + -- i -> p + -- \ / + -- r r + r <- VGM.exchange rSeq (coerce i) p + VGM.write lSeq (coerce p) r + pure r + else do + -- p i + -- \ / + -- i -> p + -- / \ + -- l l + l <- VGM.exchange lSeq (coerce i) p + VGM.write rSeq (coerce p) l + pure l + + pp <- VGM.read pSeq $ coerce p + unless (P.nullIndex pp) $ do + -- pp pp + -- / -> / + -- p i + VGM.modify lSeq (\ppl -> if ppl == p then i else ppl) $ coerce pp + -- pp pp + -- \ -> \ + -- p i + VGM.modify rSeq (\ppr -> if ppr == p then i else ppr) $ coerce pp + + -- set parents + VGM.write pSeq (coerce i) pp + VGM.write pSeq (coerce p) i + unless (P.nullIndex c) $ do + VGM.write pSeq (coerce c) p + +-- | Amortized \(O(\log n)\). Moves up a node to be a root. +-- +-- @since 1.2.0.0 +{-# INLINEABLE splayST #-} +splayST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Bool -> ST s () +splayST seq@Seq {..} i doneParentProp = do + if doneParentProp + then propNodeST seq i + else propNodeFromRootST seq i + + let inner = do + p <- VGM.read pSeq $ coerce i + unless (P.nullIndex p) $ do + pp <- VGM.read pSeq $ coerce p + if P.nullIndex pp + then do + rotateST seq i + updateNodeST seq p + pure () + else do + pl <- VGM.read lSeq $ coerce p + pr <- VGM.read rSeq $ coerce p + ppl <- VGM.read lSeq $ coerce pp + ppr <- VGM.read rSeq $ coerce pp + if pl == i && ppl == p || pr == i && ppr == p + then do + -- same direction twice + rotateST seq p + rotateST seq i + else do + rotateST seq i + rotateST seq i + updateNodeST seq pp + updateNodeST seq p + inner + + inner + updateNodeST seq i + +-- | Amortized \(O(\log n)\). Finds \(k\)-th node and splays it. Returns the new root. +-- +-- ==== Constraints +-- - \(0 \le k \lt n\) +-- +-- @since 1.2.0.0 +{-# INLINEABLE splayKthST #-} +splayKthST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s P.Index +splayKthST seq@Seq {..} root0 k0 = do + size <- VGM.read sSeq $ coerce root0 + let !_ = ACIA.checkIndex "AtCoder.Extra.Seq.Raw.splayKthST" k0 size + + let inner root k = do + propNodeST seq root + l <- VGM.read lSeq $ coerce root + -- The number of left children = the node's index counting from the leftmost. + sizeL <- if P.nullIndex l then pure 0 else VGM.read sSeq $ coerce l + case compare k sizeL of + EQ -> pure root + LT -> inner l k + GT -> do + r <- VGM.read rSeq $ coerce root + inner r (k - (sizeL + 1)) + + target <- inner root0 k0 + splayST seq target True + pure target + +-- ------------------------------------------------------------------------------------------------- +-- Bisection methods +-- ------------------------------------------------------------------------------------------------- + +-- | Amortized \(O(\log n)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINE ilowerBoundST #-} +ilowerBoundST :: + (SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq s f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> Bool) -> + -- | (r, root) + ST s (Int, P.Index) +ilowerBoundST seq root f = stToPrim $ do + (!r, !_, !root') <- imaxRightST seq root f + splayST seq root' True + pure (r, root') + +-- | Amortized \(O(\log n)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINE ilowerBoundM #-} +ilowerBoundM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> m Bool) -> + -- | (r, root) + m (Int, P.Index) +ilowerBoundM seq root f = do + (!r, !_, !root') <- imaxRightM seq root f + stToPrim $ splayST seq root' True + pure (r, root') + +-- | Amortized \(O(\log n)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINE ilowerBoundProdST #-} +ilowerBoundProdST :: + (SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq s f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid product + (Int -> a -> Bool) -> + -- | (r, root) + ST s (Int, P.Index) +ilowerBoundProdST seq root f = do + (!r, !_, !root') <- imaxRightProdST seq root f + pure (r, root') + +-- | Amortized \(O(\log n)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINE ilowerBoundProdM #-} +ilowerBoundProdM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid product + (Int -> a -> m Bool) -> + -- | (r, root) + m (Int, P.Index) +ilowerBoundProdM seq root f = do + (!r, !_, !root') <- imaxRightProdM seq root f + pure (r, root') + +-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\) +-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINE isplitMaxRightST #-} +isplitMaxRightST :: + (SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq s f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> Bool) -> + -- | (left, right) sequences where \(f\) holds for the left + ST s (P.Index, P.Index) +isplitMaxRightST seq root f = stToPrim $ isplitMaxRightM seq root (\i x -> pure (f i x)) + +-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\) +-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINEABLE isplitMaxRightM #-} +isplitMaxRightM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> m Bool) -> + -- | (left, right) sequences where \(f\) holds for the left + m (P.Index, P.Index) +isplitMaxRightM seq@Seq {..} root f + | P.nullIndex root = pure (P.undefIndex, P.undefIndex) + | otherwise = do + stToPrim $ assertRootST seq root + (!_, !c, !_) <- imaxRightM seq root f + if P.nullIndex c + then stToPrim $ do + -- `f` does hot hold + splayST seq root True + pure (P.undefIndex, root) + else stToPrim $ do + splayST seq c True + right <- VGM.read rSeq (coerce c) + if P.nullIndex right + then do + -- `f` holds for the whole sequence + pure (c, P.undefIndex) + else do + -- `f` holds for part of the sequence. detach the right child + VGM.write pSeq (coerce right) P.undefIndex + VGM.write rSeq (coerce c) P.undefIndex + updateNodeST seq c + pure (c, right) + +-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\) +-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINE isplitMaxRightProdST #-} +isplitMaxRightProdST :: + (SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq s f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid value + (Int -> a -> Bool) -> + -- | (left, right) sequences where \(f\) holds for the left + ST s (P.Index, P.Index) +isplitMaxRightProdST seq root f = stToPrim $ isplitMaxRightProdM seq root (\i x -> pure (f i x)) + +-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\) +-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINEABLE isplitMaxRightProdM #-} +isplitMaxRightProdM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + -- | \(r\) + (Int -> a -> m Bool) -> + -- | (left, right) sequences where \(f\) holds for the left + m (P.Index, P.Index) +isplitMaxRightProdM seq@Seq {..} root f + | P.nullIndex root = pure (P.undefIndex, P.undefIndex) + | otherwise = do + stToPrim $ assertRootST seq root + (!_, !c, !_) <- imaxRightProdM seq root f + if P.nullIndex c + then stToPrim $ do + -- `f` does hot hold + splayST seq root True + pure (P.undefIndex, root) + else stToPrim $ do + splayST seq c True + right <- VGM.read rSeq (coerce c) + if P.nullIndex right + then do + -- `f` holds for the whole sequence + pure (c, P.undefIndex) + else do + -- `f` holds for part of the sequence. detach the right child + VGM.write pSeq (coerce right) P.undefIndex + VGM.write rSeq (coerce c) P.undefIndex + updateNodeST seq c + pure (c, right) + +-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v\) +-- where \(f(v)\) holds for every \(v_i (0 \le i \lt k)\). Note that \(f\) works for a single +-- node, not a monoid product. +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINE imaxRightST #-} +imaxRightST :: + (SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq s f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> Bool) -> + -- | (r, left, right) + ST s (Int, P.Index, P.Index) +imaxRightST seq root0 f = stToPrim $ imaxRightM seq root0 (\i x -> pure (f i x)) + +-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\) +-- where \(f(v)\) holds for every \(v_i (0 \le i \le k)\). Note that \(f\) works for a single +-- node, not a monoid product. +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINEABLE imaxRightM #-} +imaxRightM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_i)\) that takes the index and the monoid value + (Int -> a -> m Bool) -> + -- | (r, left, right) + m (Int, P.Index, P.Index) +imaxRightM seq@Seq {..} root0 f = do + let inner offset parent root lastYes + | P.nullIndex root = pure (offset, lastYes, parent) + | otherwise = do + stToPrim $ propNodeST seq root + l <- stToPrim $ VGM.read lSeq (coerce root) + v <- stToPrim $ VGM.read vSeq (coerce root) + pos <- stToPrim $ do + if P.nullIndex l + then pure offset + else (offset +) <$> VGM.read sSeq (coerce l) + b <- f pos v + if b + then do + r <- stToPrim $ VGM.read rSeq $ coerce root + inner (pos + 1) root r root + else do + inner offset root l lastYes + + (!r, !yes, !root') <- inner 0 P.undefIndex root0 P.undefIndex + stToPrim $ splayST seq root' True + pure (r, yes, root') + +-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\) +-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINE imaxRightProdST #-} +imaxRightProdST :: + (SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq s f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid value + (Int -> a -> Bool) -> + -- | (ilowerBound, rightmost node, new root) + ST s (Int, P.Index, P.Index) +imaxRightProdST seq root0 f = imaxRightProdM seq root0 (\i x -> pure (f i x)) + +-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\) +-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\). +-- +-- ==== Constraints +-- - The node must be a root. +-- +-- @since 1.2.0.0 +{-# INLINEABLE imaxRightProdM #-} +imaxRightProdM :: + (PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | Sequence storage + Seq (PrimState m) f a -> + -- | Root node + P.Index -> + -- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid value + (Int -> a -> m Bool) -> + -- | (ilowerBound, rightmost node, new root) + m (Int, P.Index, P.Index) +imaxRightProdM seq@Seq {..} root0 f = do + let inner !acc offset parent root lastYes + | P.nullIndex root = pure (offset, lastYes, parent) + | otherwise = do + stToPrim $ propNodeST seq root + l <- stToPrim $ VGM.read lSeq $ coerce root + pos <- stToPrim $ do + if P.nullIndex l + then pure offset + else (offset +) <$> VGM.read sSeq (coerce l) + -- [0, pos] + prodM <- stToPrim $ do + -- detach right child (temporarily) and read the product + rootR <- VGM.exchange rSeq (coerce root) P.undefIndex + updateNodeST seq root + prodRoot <- VGM.read prodSeq (coerce root) + -- attach the right child again + VGM.write rSeq (coerce root) rootR + updateNodeST seq root + pure $! acc <> prodRoot + b <- f pos prodM + if b + then do + r <- stToPrim $ VGM.read rSeq $ coerce root + inner prodM (pos + 1) root r root + else do + inner acc offset root l lastYes + + (!r, !yes, !root') <- inner mempty 0 P.undefIndex root0 P.undefIndex + stToPrim $ splayST seq root' True + pure (r, yes, root') + +-- ------------------------------------------------------------------------------------------------- +-- Conversions +-- ------------------------------------------------------------------------------------------------- + +-- | Amortized \(O(n)\). Returns the sequence of monoid values. +-- +-- @since 1.2.0.0 +{-# INLINE freezeST #-} +freezeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s (VU.Vector a) +freezeST seq@Seq {sSeq, lSeq, rSeq, vSeq} root0 = do + size <- VGM.read sSeq (coerce root0) + res <- VUM.unsafeNew size + let inner i root + | P.nullIndex root = pure i + | otherwise = do + -- visit from left to right + propNodeST seq root + i' <- inner i =<< VGM.read lSeq (coerce root) + vx <- VGM.read vSeq (coerce root) + VGM.write res i' vx + inner (i' + 1) =<< VGM.read rSeq (coerce root) + _ <- inner 0 root0 + VU.unsafeFreeze res + +-- ------------------------------------------------------------------------------------------------- +-- Node methods +-- ------------------------------------------------------------------------------------------------- + +-- NOTE(pref): inlining these functions are important for the speed + +-- | \(O(1)\) Recomputes the node size and the monoid product. +{-# INLINEABLE updateNodeST #-} +updateNodeST :: (Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s () +updateNodeST Seq {..} i = do + l <- VGM.read lSeq (coerce i) + r <- VGM.read rSeq (coerce i) + prodM <- VGM.read vSeq (coerce i) + (!size', !prod') <- + if P.nullIndex l + then pure (1, prodM) + else do + sizeL <- VGM.read sSeq (coerce l) + prodL <- VGM.read prodSeq (coerce l) + pure (sizeL + 1, prodL <> prodM) + (!size'', !prod'') <- + if P.nullIndex r + then pure (size', prod') + else do + sizeR <- VGM.read sSeq (coerce r) + prodR <- VGM.read prodSeq (coerce r) + pure (size' + sizeR, prod' <> prodR) + VGM.write sSeq (coerce i) size'' + VGM.write prodSeq (coerce i) prod'' + +-- | \(O(1)\) Writes to the monoid. +{-# INLINE writeNodeST #-} +writeNodeST :: (Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> a -> ST s () +writeNodeST seq@Seq {..} root v = do + assertRootST seq root + VGM.write vSeq (coerce root) v + updateNodeST seq root + +-- | \(O(1)\) Modifies the monoid. +{-# INLINE modifyNodeST #-} +modifyNodeST :: (HasCallStack, Monoid a, VU.Unbox a) => Seq s f a -> (a -> a) -> P.Index -> ST s () +modifyNodeST seq@Seq {..} f root = do + assertRootST seq root + VGM.modify vSeq f $ coerce root + updateNodeST seq root + +-- | \(O(1)\) Modifies the monoid. +{-# INLINE exchangeNodeST #-} +exchangeNodeST :: (HasCallStack, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> a -> ST s a +exchangeNodeST seq@Seq {..} root v = do + assertRootST seq root + res <- VGM.exchange vSeq (coerce root) v + updateNodeST seq root + pure res + +-- | \(O(1)\) Swaps the left and the right children. +{-# INLINE swapLrNodeST #-} +swapLrNodeST :: Seq s f a -> P.Index -> ST s () +swapLrNodeST Seq {..} i = do + VGM.modifyM lSeq (VGM.exchange rSeq (coerce i)) (coerce i) + +-- | \(O(1)\) Reverses the left and the right children, lazily and recursively. +{-# INLINE reverseNodeST #-} +reverseNodeST :: Seq s f a -> P.Index -> ST s () +reverseNodeST seq@Seq {..} i = do + swapLrNodeST seq i + -- lazily propagate new reverse or cancel: + VGM.modify revSeq (xor (Bit True)) $ coerce i + +-- | Amortized \(O(\log n)\). Propgates the lazily propagated values on a node. +{-# INLINE propNodeST #-} +-- NOTE(pref): Although this function is large, inlining it needs for the speed. +propNodeST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s () +propNodeST seq@Seq {..} i = do + -- action + act <- VGM.exchange lazySeq (coerce i) mempty + when (act /= mempty) $ do + l <- VGM.read lSeq $ coerce i + unless (P.nullIndex l) $ do + applyNodeST seq l act + r <- VGM.read rSeq $ coerce i + unless (P.nullIndex r) $ do + applyNodeST seq r act + + -- reverse + Bit b <- VGM.exchange revSeq (coerce i) (Bit False) + when b $ do + l <- VGM.read lSeq $ coerce i + unless (P.nullIndex l) $ do + -- propagate new reverse or cancel: + reverseNodeST seq l + r <- VGM.read rSeq $ coerce i + unless (P.nullIndex r) $ do + -- propagate new reverse or cancel: + reverseNodeST seq r + +-- | Amortized \(O(\log n)\). Propagetes from the root to the given node. +{-# INLINE propNodeFromRootST #-} +propNodeFromRootST :: (HasCallStack, SegAct f a, VU.Unbox f, VU.Unbox a, Monoid a) => Seq s f a -> P.Index -> ST s () +propNodeFromRootST Seq {..} i0 = inner i0 + where + inner i = do + p <- VGM.read pSeq $ coerce i + unless (P.nullIndex p) $ do + inner p + inner i + +-- | Amortized \(O(\log n)\). Propgates at a node. +{-# INLINE applyNodeST #-} +applyNodeST :: (HasCallStack, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> f -> ST s () +applyNodeST Seq {..} i act = do + len <- VGM.read sSeq $ coerce i + VGM.modify vSeq (segAct act) $ coerce i + VGM.modify prodSeq (segActWithLength len act) $ coerce i + VGM.modify lazySeq (act <>) $ coerce i diff --git a/src/AtCoder/Extra/Tree/Lct.hs b/src/AtCoder/Extra/Tree/Lct.hs --- a/src/AtCoder/Extra/Tree/Lct.hs +++ b/src/AtCoder/Extra/Tree/Lct.hs @@ -133,7 +133,7 @@ data Lct s a = Lct { -- | The number of vertices. -- --- @since 1.1.1.0 + -- @since 1.1.1.0 nLct :: {-# UNPACK #-} !Int, -- | Decomposed node data storage: left children. -- @@ -463,7 +463,7 @@ -- | \(O(1)\) Called on changing a path-parent edge. This is for subtree folding. {-# INLINEABLE changeLightST #-} changeLightST :: Lct s a -> Vertex -> Vertex -> Vertex -> ST s () -changeLightST lct u v p = do +changeLightST _lct _u _v _p = do pure () -- | \(O(1)\) Called on erasing a path-parent edge. This is for subtree folding. diff --git a/src/AtCoder/Extra/WaveletMatrix/Raw.hs b/src/AtCoder/Extra/WaveletMatrix/Raw.hs --- a/src/AtCoder/Extra/WaveletMatrix/Raw.hs +++ b/src/AtCoder/Extra/WaveletMatrix/Raw.hs @@ -100,7 +100,7 @@ -- | \(O(n \log n)\) Creates a `RawWaveletMatrix` from a vector \(a\). -- -- @since 1.1.0.0 -{-# INLINE build #-} +{-# INLINEABLE build #-} build :: (HasCallStack) => -- | The number of different values in the compressed vector. @@ -157,7 +157,7 @@ -- original array if you can. -- -- @since 1.1.0.0 -{-# INLINABLE access #-} +{-# INLINEABLE access #-} access :: RawWaveletMatrix -> Int -> Maybe Int access RawWaveletMatrix {..} i0 | ACIA.testIndex i0 lengthRwm = @@ -181,7 +181,7 @@ -- | \(O(\log |A|)\) Goes down the wavelet matrix for collecting the kth smallest value. -- -- @since 1.1.0.0 -{-# INLINABLE goDown #-} +{-# INLINEABLE goDown #-} goDown :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int, Int, Int) goDown RawWaveletMatrix {..} l_ r_ k_ = V.ifoldl' step (0 :: Int, l_, r_, k_) bitsRwm where @@ -207,7 +207,7 @@ -- | \(O(\log |A|)\) Goes up the wavelet matrix for collecting the value \(x\). -- -- @since 1.1.0.0 -{-# INLINABLE goUp #-} +{-# INLINEABLE goUp #-} goUp :: RawWaveletMatrix -> Int -> Int -> Maybe Int goUp RawWaveletMatrix {..} i0 x = V.ifoldM' @@ -222,7 +222,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [0, y_0)\). -- -- @since 1.1.0.0 -{-# INLINABLE rankLT #-} +{-# INLINEABLE rankLT #-} rankLT :: RawWaveletMatrix -> Int -> Int -> Int -> Int rankLT RawWaveletMatrix {..} l_ r_ xr -- REMARK: This is required. The function below cannot handle the case N = 2^i and xr = N. @@ -248,7 +248,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\). -- -- @since 1.1.0.0 -{-# INLINABLE rank #-} +{-# INLINEABLE rank #-} rank :: RawWaveletMatrix -> -- | \(l\) @@ -264,7 +264,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\). -- -- @since 1.1.0.0 -{-# INLINABLE rankBetween #-} +{-# INLINEABLE rankBetween #-} rankBetween :: RawWaveletMatrix -> -- | \(l\) @@ -283,7 +283,7 @@ -- not found. -- -- @since 1.1.0.0 -{-# INLINABLE select #-} +{-# INLINEABLE select #-} select :: RawWaveletMatrix -> Int -> Maybe Int select wm = selectKth wm 0 @@ -291,7 +291,7 @@ -- if no such occurrence exists. -- -- @since 1.1.0.0 -{-# INLINABLE selectKth #-} +{-# INLINEABLE selectKth #-} selectKth :: RawWaveletMatrix -> -- | \(k\) @@ -306,7 +306,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0 -{-# INLINABLE selectIn #-} +{-# INLINEABLE selectIn #-} selectIn :: -- | A wavelet matrix RawWaveletMatrix -> @@ -324,7 +324,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0 -{-# INLINABLE selectKthIn #-} +{-# INLINEABLE selectKthIn #-} selectKthIn :: RawWaveletMatrix -> -- | \(l\) @@ -369,7 +369,7 @@ -- largest value. Note that duplicated values are counted as distinct occurrences. -- -- @since 1.1.0.0 -{-# INLINABLE kthLargestIn #-} +{-# INLINEABLE kthLargestIn #-} kthLargestIn :: -- | A wavelet matrix RawWaveletMatrix -> @@ -390,7 +390,7 @@ -- \(k\)-th (0-based) largest value. Note that duplicated values are counted as distinct occurrences. -- -- @since 1.1.0.0 -{-# INLINABLE ikthLargestIn #-} +{-# INLINEABLE ikthLargestIn #-} ikthLargestIn :: -- | A wavelet matrix RawWaveletMatrix -> @@ -411,7 +411,7 @@ -- smallest value. Note that duplicated values are counted as distinct occurrences. -- -- @since 1.1.0.0 -{-# INLINABLE kthSmallestIn #-} +{-# INLINEABLE kthSmallestIn #-} kthSmallestIn :: -- | A wavelet matrix RawWaveletMatrix -> @@ -432,7 +432,7 @@ -- \(k\)-th (0-based) smallest value. Note that duplicated values are counted as distinct occurrences. -- -- @since 1.1.0.0 -{-# INLINABLE ikthSmallestIn #-} +{-# INLINEABLE ikthSmallestIn #-} ikthSmallestIn :: RawWaveletMatrix -> -- | \(l\) @@ -452,21 +452,21 @@ -- values are counted as distinct occurrences. -- -- @since 1.1.0.0 -{-# INLINABLE unsafeKthLargestIn #-} +{-# INLINEABLE 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 -{-# INLINABLE unsafeIKthLargestIn #-} +{-# INLINEABLE 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 -{-# INLINABLE unsafeKthSmallestIn #-} +{-# INLINEABLE unsafeKthSmallestIn #-} unsafeKthSmallestIn :: RawWaveletMatrix -> Int -> Int -> Int -> Int unsafeKthSmallestIn wm l_ r_ k_ = let (!x, !_, !_, !_) = goDown wm l_ r_ k_ @@ -475,7 +475,7 @@ -- | \(O(\log a)\) -- -- @since 1.1.0.0 -{-# INLINABLE unsafeIKthSmallestIn #-} +{-# INLINEABLE unsafeIKthSmallestIn #-} unsafeIKthSmallestIn :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int) unsafeIKthSmallestIn wm l_ r_ k_ = let (!x, !l, !_, !k) = goDown wm l_ r_ k_ @@ -485,7 +485,7 @@ -- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0]\). -- -- @since 1.1.0.0 -{-# INLINABLE lookupLE #-} +{-# INLINEABLE lookupLE #-} lookupLE :: -- | A wavelet matrix RawWaveletMatrix -> @@ -510,7 +510,7 @@ -- | \(O(\log a)\) Finds the maximum \(x\) in \([l, r)\) s.t. \(x_{0} \lt x\). -- -- @since 1.1.0.0 -{-# INLINABLE lookupLT #-} +{-# INLINEABLE lookupLT #-} lookupLT :: RawWaveletMatrix -> -- | \(l\) @@ -526,7 +526,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times [y_0, \infty)\). -- -- @since 1.1.0.0 -{-# INLINABLE lookupGE #-} +{-# INLINEABLE lookupGE #-} lookupGE :: RawWaveletMatrix -> -- | \(l\) @@ -551,7 +551,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\). -- -- @since 1.1.0.0 -{-# INLINABLE lookupGT #-} +{-# INLINEABLE lookupGT #-} lookupGT :: RawWaveletMatrix -> -- | \(l\) @@ -568,14 +568,14 @@ -- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0 -{-# INLINABLE assocsIn #-} +{-# INLINEABLE 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 -{-# INLINABLE assocsWith #-} +{-# INLINEABLE assocsWith #-} assocsWith :: RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)] assocsWith RawWaveletMatrix {..} l_ r_ f | l'_ < r'_ = inner (0 :: Int) (0 :: Int) l'_ r'_ [] @@ -611,14 +611,14 @@ -- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0 -{-# INLINABLE descAssocsIn #-} +{-# INLINEABLE 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 -{-# INLINABLE descAssocsInWith #-} +{-# INLINEABLE descAssocsInWith #-} descAssocsInWith :: RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)] descAssocsInWith RawWaveletMatrix {..} l_ r_ f | l'_ < r'_ = inner (0 :: Int) (0 :: Int) l'_ r'_ [] diff --git a/src/AtCoder/Extra/WaveletMatrix2d.hs b/src/AtCoder/Extra/WaveletMatrix2d.hs --- a/src/AtCoder/Extra/WaveletMatrix2d.hs +++ b/src/AtCoder/Extra/WaveletMatrix2d.hs @@ -35,6 +35,8 @@ -- >>> 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} +-- +-- @since 1.1.0.0 module AtCoder.Extra.WaveletMatrix2d ( -- * Wavelet matrix 2D WaveletMatrix2d (..), @@ -59,7 +61,8 @@ 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 Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) +import Control.Monad.ST (ST) import Data.Bit (Bit (..)) import Data.Bits (Bits (testBit)) import Data.Maybe (fromJust, fromMaybe) @@ -77,22 +80,36 @@ -- - `maxRight` can be implemented. -- | Segment Tree on Wavelet Matrix: points on a 2D plane and rectangle products. +-- +-- @since 1.1.0.0 data WaveletMatrix2d s a = WaveletMatrix2d { -- | The wavelet matrix that represents points on a 2D plane. + -- + -- @since 1.1.0.0 rawWmWm2d :: !Rwm.RawWaveletMatrix, -- | (x, y) index compression dictionary. + -- + -- @since 1.1.0.0 xyDictWm2d :: !(VU.Vector (Int, Int)), -- | y index compression dictionary. + -- + -- @since 1.1.0.0 yDictWm2d :: !(VU.Vector Int), -- | The segment tree of the weights of the points in the order of `xyDictWm2d`. + -- + -- @since 1.1.0.0 segTreesWm2d :: !(V.Vector (ST.SegTree s a)), -- | The inverse operator of the interested monoid. + -- + -- @since 1.1.0.0 invWm2d :: !(a -> a) } -- | \(O(n \log n)\) Creates a `WaveletMatrix2d` with `mempty` as the initial monoid -- values for each point. -{-# INLINE new #-} +-- +-- @since 1.1.0.0 +{-# INLINEABLE new #-} new :: (PrimMonad m, Monoid a, VU.Unbox a) => -- | Inverse operator of the monoid @@ -101,7 +118,7 @@ VU.Vector (Int, Int) -> -- | A 2D wavelet matrix m (WaveletMatrix2d (PrimState m) a) -new invWm2d xys = do +new invWm2d xys = stToPrim $ do let n = VG.length xys let xyDictWm2d = VU.uniq . VU.modify (VAI.sortBy compare) $ xys let (!_, !ys) = VU.unzip xys @@ -115,7 +132,9 @@ -- | \(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 #-} +-- +-- @since 1.1.0.0 +{-# INLINEABLE build #-} build :: (PrimMonad m, Monoid a, VU.Unbox a) => -- | Inverse operator of the monoid @@ -124,7 +143,7 @@ VU.Vector (Int, Int, a) -> -- | A 2D wavelet matrix m (WaveletMatrix2d (PrimState m) a) -build invWm2d xysw = do +build invWm2d xysw = stToPrim $ do let (!xs, !ys, !_) = VU.unzip3 xysw wm <- new invWm2d $ VU.zip xs ys -- not the fastest implementation though @@ -133,15 +152,19 @@ 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 +-- +-- @since 1.1.0.0 +{-# INLINEABLE read #-} +read :: (HasCallStack, PrimMonad m, VU.Unbox a, Monoid a) => 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 +-- +-- @since 1.1.0.0 +{-# INLINEABLE write #-} +write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => WaveletMatrix2d (PrimState m) a -> (Int, Int) -> a -> m () +write WaveletMatrix2d {..} (!x, !y) v = stToPrim $ do let !i_ = fromJust $ lowerBound xyDictWm2d (x, y) V.ifoldM'_ ( \i iRow (!bits, !seg) -> do @@ -158,9 +181,11 @@ -- | \(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 +-- +-- @since 1.1.0.0 +{-# INLINEABLE modify #-} +modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => WaveletMatrix2d (PrimState m) a -> (a -> a) -> (Int, Int) -> m () +modify WaveletMatrix2d {..} f (!x, !y) = stToPrim $ do let !i_ = fromJust $ lowerBound xyDictWm2d (x, y) V.ifoldM'_ ( \i iRow (!bits, !seg) -> do @@ -176,8 +201,10 @@ $ 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 +-- +-- @since 1.1.0.0 +{-# INLINEABLE prod #-} +prod :: (HasCallStack, PrimMonad m, VU.Unbox a, Monoid a) => 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' @@ -193,8 +220,10 @@ -- | \(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) +-- +-- @since 1.1.0.0 +{-# INLINEABLE prodMaybe #-} +prodMaybe :: (PrimMonad m, VU.Unbox a, Monoid a) => 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 @@ -209,23 +238,27 @@ 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 +-- +-- @since 1.1.0.0 +{-# INLINEABLE allProd #-} +allProd :: (HasCallStack, PrimMonad m, PrimMonad m, VU.Unbox a, Monoid a) => 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. +-- +-- @since 1.1.0.0 {-# 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 +unsafeProd :: (PrimMonad m, VU.Unbox a, Monoid a) => WaveletMatrix2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a +unsafeProd wm xl' xr' yl' yr' = stToPrim $ 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 +{-# INLINEABLE prodLT #-} +prodLT :: (Monoid a, VU.Unbox a) => WaveletMatrix2d s a -> Int -> Int -> Int -> ST s a prodLT WaveletMatrix2d {..} !l_ !r_ yUpper = do (!res, !_, !_) <- do V.ifoldM' @@ -246,32 +279,3 @@ (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/Internal/Convolution.hs b/src/AtCoder/Internal/Convolution.hs --- a/src/AtCoder/Internal/Convolution.hs +++ b/src/AtCoder/Internal/Convolution.hs @@ -26,7 +26,7 @@ import AtCoder.ModInt qualified as AM import Control.Monad (when) import Control.Monad.Fix (fix) -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (ST) import Data.Bits (bit, complement, countTrailingZeros, (.<<.), (.>>.)) import Data.Foldable import Data.Vector.Generic qualified as VG @@ -58,8 +58,8 @@ -- | \(O(\log m)\) Creates an `FftInfo`. -- -- @since 1.0.0.0 -{-# INLINE newInfo #-} -newInfo :: forall m p. (PrimMonad m, AM.Modulus p) => m (FftInfo p) +{-# INLINABLE newInfo #-} +newInfo :: forall s p. (AM.Modulus p) => ST s (FftInfo p) newInfo = do let !g = AM.primitiveRootModulus (proxy# @p) let !m = fromIntegral $ natVal' (proxy# @p) @@ -111,13 +111,13 @@ pure FftInfo {..} -- | @since 1.0.0.0 -{-# INLINE butterfly #-} +{-# INLINABLE butterfly #-} butterfly :: - forall m p. - (PrimMonad m, AM.Modulus p) => + forall s p. + (AM.Modulus p) => FftInfo p -> - VUM.MVector (PrimState m) (AM.ModInt p) -> - m () + VUM.MVector s (AM.ModInt p) -> + ST s () butterfly FftInfo {..} a = do let n = VUM.length a let h = countTrailingZeros n @@ -175,13 +175,13 @@ loop $ len + 2 -- | @since 1.0.0.0 -{-# INLINE butterflyInv #-} +{-# INLINABLE butterflyInv #-} butterflyInv :: - forall m p. - (PrimMonad m, AM.Modulus p) => + forall s p. + (AM.Modulus p) => FftInfo p -> - VUM.MVector (PrimState m) (AM.ModInt p) -> - m () + VUM.MVector s (AM.ModInt p) -> + ST s () butterflyInv FftInfo {..} a = do let n = VUM.length a let h = countTrailingZeros n @@ -240,7 +240,7 @@ loop $ len - 2 -- | @since 1.0.0.0 -{-# INLINE convolutionNaive #-} +{-# INLINABLE convolutionNaive #-} convolutionNaive :: forall p. (AM.Modulus p) => diff --git a/src/AtCoder/Internal/McfCsr.hs b/src/AtCoder/Internal/McfCsr.hs --- a/src/AtCoder/Internal/McfCsr.hs +++ b/src/AtCoder/Internal/McfCsr.hs @@ -16,7 +16,8 @@ ) where -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) +import Control.Monad.ST (ST) import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU @@ -42,12 +43,9 @@ costCsr :: !(VU.Vector cost) } --- | \(O(n + m)\) Creates `Csr`. --- --- @since 1.0.0.0 -{-# INLINE build #-} -build :: (HasCallStack, Num cap, VU.Unbox cap, VU.Unbox cost, Num cost, PrimMonad m) => Int -> VU.Vector (Int, Int, cap, cap, cost) -> m (VU.Vector Int, Csr (PrimState m) cap cost) -build n edges = do +{-# INLINEABLE buildST #-} +buildST :: (HasCallStack, Num cap, VU.Unbox cap, VU.Unbox cost, Num cost) => Int -> VU.Vector (Int, Int, cap, cap, cost) -> ST s (VU.Vector Int, Csr s cap cost) +buildST n edges = do let m = VU.length edges -- craete the offsets first (this is a different step from ac-librar) let startCsr = VU.create $ do @@ -91,6 +89,13 @@ revCsr <- VU.unsafeFreeze revVec costCsr <- VU.unsafeFreeze costVec pure (edgeIdx, Csr {..}) + +-- | \(O(n + m)\) Creates `Csr`. +-- +-- @since 1.0.0.0 +{-# INLINE build #-} +build :: (HasCallStack, PrimMonad m, Num cap, VU.Unbox cap, VU.Unbox cost, Num cost) => Int -> VU.Vector (Int, Int, cap, cap, cost) -> m (VU.Vector Int, Csr (PrimState m) cap cost) +build n edges = stToPrim $ buildST n edges -- | \(O(1)\) Returns a vector of @(to, rev, cost)@. -- 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 @@ -6,7 +6,7 @@ -- -- ==== __Example__ -- >>> import AtCoder.Internal.MinHeap qualified as MH --- >>> heap <- MH.new @Int 4 +-- >>> heap <- MH.new @_ @Int 4 -- >>> MH.capacity heap -- 4 -- @@ -54,7 +54,8 @@ where import Control.Monad (when) -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) +import Control.Monad.ST (ST) import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM @@ -85,7 +86,7 @@ -- -- @since 1.0.0.0 {-# INLINE new #-} -new :: (VU.Unbox a, PrimMonad m) => Int -> m (Heap (PrimState m) a) +new :: (PrimMonad m, VU.Unbox a) => Int -> m (Heap (PrimState m) a) new n = do sizeBH_ <- VUM.replicate 1 0 dataBH <- VUM.unsafeNew n @@ -102,29 +103,26 @@ -- -- @since 1.0.0.0 {-# INLINE length #-} -length :: (VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m Int +length :: (PrimMonad m, VU.Unbox a) => Heap (PrimState m) a -> m Int length Heap {sizeBH_} = VGM.unsafeRead sizeBH_ 0 -- | \(O(1)\) Returns `True` if the heap is empty. -- -- @since 1.0.0.0 {-# INLINE null #-} -null :: (VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m Bool +null :: (PrimMonad m, VU.Unbox a) => Heap (PrimState m) a -> m Bool null = (<$>) (== 0) . length -- | \(O(1)\) Sets the `length` to zero. -- -- @since 1.0.0.0 {-# INLINE clear #-} -clear :: (VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m () +clear :: (PrimMonad m, VU.Unbox a) => Heap (PrimState m) a -> m () clear Heap {sizeBH_} = VGM.unsafeWrite sizeBH_ 0 0 --- | \(O(\log n)\) Inserts an element to the heap. --- --- @since 1.0.0.0 -{-# INLINE push #-} -push :: (HasCallStack, Ord a, VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> a -> m () -push Heap {..} x = do +{-# INLINEABLE pushST #-} +pushST :: (HasCallStack, Ord a, VU.Unbox a) => Heap s a -> a -> ST s () +pushST Heap {..} x = do i0 <- VGM.unsafeRead sizeBH_ 0 VGM.write dataBH i0 x VGM.unsafeWrite sizeBH_ 0 $ i0 + 1 @@ -136,13 +134,16 @@ siftUp iParent siftUp i0 --- | \(O(\log n)\) Removes the last element from the heap and returns it, or `Nothing` if it is --- empty. +-- | \(O(\log n)\) Inserts an element to the heap. -- -- @since 1.0.0.0 -{-# INLINE pop #-} -pop :: (HasCallStack, Ord a, VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m (Maybe a) -pop heap@Heap {..} = do +{-# INLINE push #-} +push :: (HasCallStack, PrimMonad m, Ord a, VU.Unbox a) => Heap (PrimState m) a -> a -> m () +push heap x = stToPrim $ pushST heap x + +{-# INLINEABLE popST #-} +popST :: (HasCallStack, Ord a, VU.Unbox a) => Heap s a -> ST s (Maybe a) +popST heap@Heap {..} = do len <- length heap if len == 0 then pure Nothing @@ -178,13 +179,21 @@ siftDown 0 pure $ Just root +-- | \(O(\log n)\) Removes the last element from the heap and returns it, or `Nothing` if it is +-- empty. +-- +-- @since 1.0.0.0 +{-# INLINE pop #-} +pop :: (HasCallStack, PrimMonad m, Ord a, VU.Unbox a) => Heap (PrimState m) a -> m (Maybe a) +pop heap = stToPrim $ popST heap + -- | \(O(\log n)\) `pop` with the return value discarded. -- -- @since 1.0.0.0 {-# INLINE pop_ #-} pop_ :: (HasCallStack, Ord a, VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m () pop_ heap = do - _ <- pop heap + _ <- stToPrim $ popST heap pure () -- | \(O(1)\) Returns the smallest value in the heap, or `Nothing` if it is empty. 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 @@ -20,7 +20,7 @@ -- ** (Extra API) CSR API sccIdsCsr, - sccCsr + sccCsr, ) where @@ -28,8 +28,8 @@ import AtCoder.Internal.GrowVec qualified as ACIGV import Control.Monad (unless, when) import Control.Monad.Fix (fix) -import Control.Monad.Primitive (PrimMonad, PrimState) -import Control.Monad.ST (runST) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) +import Control.Monad.ST (ST, runST) import Data.Foldable (for_) import Data.Maybe (fromJust) import Data.Vector qualified as V @@ -75,12 +75,9 @@ 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 +-- NOTE(perf): faster without INLINEABLE (somehow) +sccST :: SccGraph s -> ST s (V.Vector (VU.Vector Int)) +sccST g = do (!groupNum, !ids) <- sccIds g let counts = VU.create $ do vec <- VUM.replicate groupNum (0 :: Int) @@ -95,10 +92,17 @@ VGM.write (groups VG.! sccId) i v V.mapM VU.unsafeFreeze groups +-- | \(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 = stToPrim $ sccST g + -- | \(O(n + m)\) API) Returns a pair of @(# of scc, scc id)@. -- -- @since 1.1.0.0 -{-# INLINABLE sccIdsCsr #-} +{-# INLINEABLE sccIdsCsr #-} sccIdsCsr :: ACICSR.Csr w -> (Int, VU.Vector Int) sccIdsCsr g@ACICSR.Csr {..} = runST $ do -- see also the Wikipedia: https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm#The_algorithm_in_pseudocode @@ -173,7 +177,7 @@ -- | \(O(n + m)\) Returns the strongly connected components. -- -- @since 1.1.0.0 -{-# INLINABLE sccCsr #-} +{-# INLINEABLE sccCsr #-} sccCsr :: ACICSR.Csr w -> V.Vector (VU.Vector Int) sccCsr g = runST $ do groups <- V.mapM VUM.unsafeNew $ VU.convert counts diff --git a/src/AtCoder/LazySegTree.hs b/src/AtCoder/LazySegTree.hs --- a/src/AtCoder/LazySegTree.hs +++ b/src/AtCoder/LazySegTree.hs @@ -122,7 +122,8 @@ import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Bit qualified as ACIBIT import Control.Monad (unless, when) -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) +import Control.Monad.ST (ST) import Data.Bits (bit, countLeadingZeros, countTrailingZeros, testBit, (.&.), (.<<.), (.>>.)) import Data.Foldable (for_) import Data.Vector.Generic.Mutable qualified as VGM @@ -289,6 +290,14 @@ segActWithLength :: Int -> f -> a -> a segActWithLength _ = segAct +-- | @since 1.2.0.0 +instance SegAct () a where + {-# INLINE segAct #-} + segAct _ = id + {-# INLINE segActWithLength #-} + segActWithLength :: Int -> f -> a -> a + segActWithLength _ _ = id + -- | A lazily propagated segment tree defined around `SegAct`. -- -- @since 1.0.0.0 @@ -311,6 +320,21 @@ lzLst :: !(VUM.MVector s f) } +{-# INLINE buildST #-} +buildST :: (Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => VU.Vector a -> ST s (LazySegTree s f a) +buildST vs = do + let nLst = VU.length vs + let sizeLst = ACIBIT.bitCeil nLst + let logLst = countTrailingZeros sizeLst + dLst <- VUM.replicate (2 * sizeLst) mempty + lzLst <- VUM.replicate sizeLst mempty + VU.iforM_ vs $ \i v -> do + VGM.write dLst (sizeLst + i) v + let segtree = LazySegTree {..} + for_ [sizeLst - 1, sizeLst - 2 .. 1] $ \i -> do + updateST segtree i + pure segtree + -- | Creates an array of length \(n\). All the elements are initialized to `mempty`. -- -- ==== Constraints @@ -323,7 +347,7 @@ {-# INLINE new #-} new :: (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Int -> m (LazySegTree (PrimState m) f a) new nLst - | nLst >= 0 = build $ VU.replicate nLst mempty + | nLst >= 0 = stToPrim $ buildST $ VU.replicate nLst mempty | otherwise = error $ "new: given negative size `" ++ show nLst ++ "`" -- | Creates an array with initial values \(vs\). @@ -337,19 +361,18 @@ -- @since 1.0.0.0 {-# INLINE build #-} build :: (PrimMonad m, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => VU.Vector a -> m (LazySegTree (PrimState m) f a) -build vs = do - let nLst = VU.length vs - let sizeLst = ACIBIT.bitCeil nLst - let logLst = countTrailingZeros sizeLst - dLst <- VUM.replicate (2 * sizeLst) mempty - lzLst <- VUM.replicate sizeLst mempty - VU.iforM_ vs $ \i v -> do - VGM.write dLst (sizeLst + i) v - let segtree = LazySegTree {..} - for_ [sizeLst - 1, sizeLst - 2 .. 1] $ \i -> do - update segtree i - pure segtree +build vs = stToPrim $ buildST vs +{-# INLINE writeST #-} +writeST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> a -> ST s () +writeST self@LazySegTree {..} p x = do + let p' = p + sizeLst + for_ [logLst, logLst - 1 .. 1] $ \i -> do + pushST self $ p' .>>. i + VGM.unsafeWrite dLst p' x + for_ [1 .. logLst] $ \i -> do + updateST self $ p' .>>. i + -- | Sets \(p\)-th value of the array to \(x\). -- -- ==== Constraints @@ -361,14 +384,19 @@ -- @since 1.0.0.0 {-# INLINE write #-} write :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> a -> m () -write self@LazySegTree {..} p x = do - let !_ = ACIA.checkIndex "AtCoder.LazySegTree.write" p nLst +write self p x = stToPrim $ writeST self p x + where + !_ = ACIA.checkIndex "AtCoder.LazySegTree.write" p (nLst self) + +{-# INLINE modifyST #-} +modifyST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> (a -> a) -> Int -> ST s () +modifyST self@LazySegTree {..} f p = do let p' = p + sizeLst for_ [logLst, logLst - 1 .. 1] $ \i -> do - push self $ p' .>>. i - VGM.write dLst p' x + pushST self $ p' .>>. i + VGM.unsafeModify dLst f p' for_ [1 .. logLst] $ \i -> do - update self $ p' .>>. i + updateST self $ p' .>>. i -- | (Extra API) Modifies \(p\)-th value with a function \(f\). -- @@ -381,14 +409,9 @@ -- @since 1.0.0.0 {-# INLINE modify #-} modify :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> (a -> a) -> Int -> m () -modify self@LazySegTree {..} f p = do - let !_ = ACIA.checkIndex "AtCoder.LazySegTree.modify" p nLst - let p' = p + sizeLst - for_ [logLst, logLst - 1 .. 1] $ \i -> do - push self $ p' .>>. i - VGM.modify dLst f p' - for_ [1 .. logLst] $ \i -> do - update self $ p' .>>. i +modify self f p = stToPrim $ modifyST self f p + where + !_ = ACIA.checkIndex "AtCoder.LazySegTree.modify" p (nLst self) -- | (Extra API) Modifies \(p\)-th value with a monadic function \(f\). -- @@ -402,13 +425,24 @@ {-# INLINE modifyM #-} modifyM :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> (a -> m a) -> Int -> m () modifyM self@LazySegTree {..} f p = do - let !_ = ACIA.checkIndex "AtCoder.LazySegTree.modify" p nLst + let !_ = ACIA.checkIndex "AtCoder.LazySegTree.modifyM" p nLst let p' = p + sizeLst - for_ [logLst, logLst - 1 .. 1] $ \i -> do - push self $ p' .>>. i + stToPrim $ for_ [logLst, logLst - 1 .. 1] $ \i -> do + pushST self $ p' .>>. i VGM.modifyM dLst f p' + stToPrim $ for_ [1 .. logLst] $ \i -> do + updateST self $ p' .>>. i + +{-# INLINE exchangeST #-} +exchangeST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> a -> ST s a +exchangeST self@LazySegTree {..} p x = do + let p' = p + sizeLst + for_ [logLst, logLst - 1 .. 1] $ \i -> do + pushST self $ p' .>>. i + res <- VGM.unsafeExchange dLst p' x for_ [1 .. logLst] $ \i -> do - update self $ p' .>>. i + updateST self $ p' .>>. i + pure res -- | (Extra API) Sets \(p\)-th value of the array to \(x\) and returns the old value. -- @@ -421,15 +455,17 @@ -- @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 +exchange self p x = stToPrim $ exchangeST self p x + where + !_ = ACIA.checkIndex "AtCoder.LazySegTree.exchange" p (nLst self) + +{-# INLINE readST #-} +readST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> ST s a +readST self@LazySegTree {..} p = do 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 + pushST self $ p' .>>. i + VGM.unsafeRead dLst p' -- | Returns \(p\)-th value of the array. -- @@ -442,12 +478,9 @@ -- @since 1.0.0.0 {-# INLINE read #-} read :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> m a -read self@LazySegTree {..} p = do - let !_ = ACIA.checkIndex "AtCoder.LazySegTree.read" p nLst - let p' = p + sizeLst - for_ [logLst, logLst - 1 .. 1] $ \i -> do - push self $ p' .>>. i - VGM.read dLst p' +read self p = stToPrim $ readST self p + where + !_ = ACIA.checkIndex "AtCoder.LazySegTree.read" p (nLst self) -- | Returns the product of \([a[l], ..., a[r - 1]]\), assuming the properties of the monoid. It -- returns `mempty` if \(l = r\). @@ -464,7 +497,7 @@ prod self@LazySegTree {nLst} l0 r0 | not (ACIA.testInterval l0 r0 nLst) = ACIA.errorInterval "AtCoder.LazySegTree.prod" l0 r0 nLst | l0 == r0 = pure mempty - | otherwise = unsafeProd self l0 r0 + | otherwise = stToPrim $ unsafeProdST self l0 r0 -- | 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 @@ -479,17 +512,17 @@ prodMaybe self@LazySegTree {nLst} l0 r0 | not (ACIA.testInterval l0 r0 nLst) = pure Nothing | l0 == r0 = pure (Just mempty) - | otherwise = Just <$> unsafeProd self l0 r0 + | otherwise = stToPrim $ Just <$> unsafeProdST self l0 r0 -- | Internal implementation of `prod`. -{-# INLINE unsafeProd #-} -unsafeProd :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> Int -> m a -unsafeProd self@LazySegTree {..} l0 r0 = do +{-# INLINE unsafeProdST #-} +unsafeProdST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> Int -> ST s a +unsafeProdST self@LazySegTree {..} l0 r0 = do let l = l0 + sizeLst let r = r0 + sizeLst for_ [logLst, logLst - 1 .. 1] $ \i -> do - when (((l .>>. i) .<<. i) /= l) $ push self $ l .>>. i - when (((r .>>. i) .<<. i) /= r) $ push self $ (r - 1) .>>. i + when (((l .>>. i) .<<. i) /= l) $ pushST self $ l .>>. i + when (((r .>>. i) .<<. i) /= r) $ pushST self $ (r - 1) .>>. i inner l (r - 1) mempty mempty where -- NOTE: we're using inclusive range [l, r] for simplicity @@ -517,66 +550,76 @@ 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\). --- --- ==== Constraints --- - \(0 \leq p \lt n\) --- --- ==== Complexity --- - \(O(\log n)\) --- --- @since 1.0.0.0 -{-# INLINE applyAt #-} -applyAt :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> f -> m () -applyAt self@LazySegTree {..} p f = do - let !_ = ACIA.checkIndex "AtCoder.LazySegTree.applyAt" p nLst +{-# INLINE applyAtST #-} +applyAtST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> f -> ST s () +applyAtST self@LazySegTree {..} p f = do let p' = p + sizeLst -- propagate for_ [logLst, logLst - 1 .. 1] $ \i -> do - push self $ p' .>>. i + pushST self $ p' .>>. i let !len = bit $! logLst - (63 - countLeadingZeros p') VGM.modify dLst (segActWithLength len f) p' -- evaluate for_ [1 .. logLst] $ \i -> do - update self $ p' .>>. i + updateST self $ p' .>>. i --- | Applies @segAct f@ to an interval \([l, r)\). +-- | Applies @segAct f@ to an index \(p\). -- -- ==== Constraints --- - \(0 \leq l \leq r \leq n\) +-- - \(0 \leq p \lt n\) -- -- ==== Complexity -- - \(O(\log n)\) -- -- @since 1.0.0.0 -{-# INLINE applyIn #-} -applyIn :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> Int -> f -> m () -applyIn self@LazySegTree {..} l0 r0 f +{-# INLINE applyAt #-} +applyAt :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> f -> m () +applyAt self p f = stToPrim $ applyAtST self p f + where + !_ = ACIA.checkIndex "AtCoder.LazySegTree.applyAt" p (nLst self) + +{-# INLINE applyInST #-} +applyInST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> Int -> f -> ST s () +applyInST self@LazySegTree {..} l0 r0 f | l0 == r0 = pure () | otherwise = do let l = l0 + sizeLst let r = r0 + sizeLst -- propagate for_ [logLst, logLst - 1 .. 1] $ \i -> do - when (((l .>>. i) .<<. i) /= l) $ push self (l .>>. i) - when (((r .>>. i) .<<. i) /= r) $ push self ((r - 1) .>>. i) + when (((l .>>. i) .<<. i) /= l) $ pushST self (l .>>. i) + when (((r .>>. i) .<<. i) /= r) $ pushST self ((r - 1) .>>. i) inner l (r - 1) -- evaluate for_ [1 .. logLst] $ \i -> do - when (((l .>>. i) .<<. i) /= l) $ update self (l .>>. i) - when (((r .>>. i) .<<. i) /= r) $ update self ((r - 1) .>>. i) + when (((l .>>. i) .<<. i) /= l) $ updateST self (l .>>. i) + when (((r .>>. i) .<<. i) /= r) $ updateST self ((r - 1) .>>. i) where - !_ = ACIA.checkInterval "AtCoder.LazySegTree.applyIn" l0 r0 nLst -- NOTE: we're using inclusive range [l, r] for simplicity inner l r | l > r = pure () | otherwise = do when (testBit l 0) $ do - allApply self l f + allApplyST self l f unless (testBit r 0) $ do - allApply self r f + allApplyST self r f inner ((l + 1) .>>. 1) ((r - 1) .>>. 1) +-- | Applies @segAct f@ to an interval \([l, r)\). +-- +-- ==== Constraints +-- - \(0 \leq l \leq r \leq n\) +-- +-- ==== Complexity +-- - \(O(\log n)\) +-- +-- @since 1.0.0.0 +{-# INLINE applyIn #-} +applyIn :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> Int -> f -> m () +applyIn self l0 r0 f = stToPrim $ applyInST self l0 r0 f + where + !_ = ACIA.checkInterval "AtCoder.LazySegTree.applyIn" l0 r0 (nLst self) + -- | Applies a binary search on the segment tree. It returns an index \(l\) that satisfies both of the -- following. -- @@ -621,8 +664,8 @@ then pure 0 else do let r = r0 + sizeLst - for_ [logLst, logLst - 1 .. 1] $ \i -> do - push self $ (r - 1) .>>. i + stToPrim $ for_ [logLst, logLst - 1 .. 1] $ \i -> do + pushST self $ (r - 1) .>>. i inner r mempty where -- NOTE: Not ordinary bounds check! @@ -643,9 +686,10 @@ | otherwise = r inner2 r sm | r < sizeLst = do - push self r let r' = 2 * r + 1 - !sm' <- (<> sm) <$> VGM.read dLst r' + sm' <- stToPrim $ do + pushST self r + (<> sm) <$> VGM.read dLst r' b <- g sm' if b then inner2 (r' - 1) sm' @@ -696,8 +740,8 @@ then pure nLst else do let l = l0 + sizeLst - for_ [logLst, logLst - 1 .. 1] $ \i -> do - push self (l .>>. i) + stToPrim $ for_ [logLst, logLst - 1 .. 1] $ \i -> do + pushST self (l .>>. i) inner l mempty where -- NOTE: Not ordinary bounds check! @@ -720,9 +764,10 @@ | otherwise = l inner2 l !sm | l < sizeLst = do - push self l let l' = 2 * l - !sm' <- (sm <>) <$> VGM.read dLst l' + sm' <- stToPrim $ do + pushST self l + (sm <>) <$> VGM.read dLst l' b <- g sm' if b then inner2 (l' + 1) sm' @@ -736,8 +781,8 @@ freeze :: (PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> m (VU.Vector a) freeze self@LazySegTree {..} = do -- push all (we _could_ skip some elements) - for_ [1 .. sizeLst - 1] $ \i -> do - push self i + stToPrim $ for_ [1 .. sizeLst - 1] $ \i -> do + pushST self i VU.freeze . VUM.take nLst $ VUM.drop sizeLst dLst -- | \(O(n)\) Unsafely converts a mutable vector to an immutable one without copying. The mutable @@ -748,32 +793,34 @@ unsafeFreeze :: (PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> m (VU.Vector a) unsafeFreeze self@LazySegTree {..} = do -- push all (we _could_ skip some elements) - for_ [1 .. sizeLst - 1] $ \i -> do - push self i + stToPrim $ for_ [1 .. sizeLst - 1] $ \i -> do + pushST self i VU.unsafeFreeze . VUM.take nLst $ VUM.drop sizeLst dLst +-- NOTE (perf): these functions have to be inlined after all: + -- | \(O(1)\) -{-# INLINE update #-} -update :: (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> m () -update LazySegTree {..} k = do +{-# INLINE updateST #-} +updateST :: (Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> ST s () +updateST LazySegTree {dLst} k = do opL <- VGM.read dLst $ 2 * k opR <- VGM.read dLst $ 2 * k + 1 VGM.write dLst k $! opL <> opR -- | \(O(1)\) -{-# INLINE allApply #-} -allApply :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> f -> m () -allApply LazySegTree {..} k f = do +{-# INLINE allApplyST #-} +allApplyST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> f -> ST s () +allApplyST LazySegTree {..} k f = do let !len = bit $! logLst - (63 - countLeadingZeros k) VGM.modify dLst (segActWithLength len f) k when (k < sizeLst) $ do VGM.modify lzLst (f <>) k -- | \(O(1)\) -{-# INLINE push #-} -push :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> m () -push self@LazySegTree {..} k = do +{-# INLINE pushST #-} +pushST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> ST s () +pushST self@LazySegTree {lzLst} k = do lzK <- VGM.read lzLst k - allApply self (2 * k) lzK - allApply self (2 * k + 1) lzK + allApplyST self (2 * k) lzK + allApplyST self (2 * k + 1) lzK VGM.write lzLst k mempty diff --git a/src/AtCoder/MaxFlow.hs b/src/AtCoder/MaxFlow.hs --- a/src/AtCoder/MaxFlow.hs +++ b/src/AtCoder/MaxFlow.hs @@ -64,7 +64,8 @@ import AtCoder.Internal.Queue qualified as ACIQ import Control.Monad (unless, when) import Control.Monad.Fix (fix) -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) +import Control.Monad.ST (ST) import Data.Bit (Bit (..)) import Data.Primitive.MutVar (readMutVar) import Data.Vector qualified as V @@ -194,7 +195,7 @@ cap -> -- | Max flow m cap -flow MfGraph {..} s t flowLimit = do +flow MfGraph {..} s t flowLimit = stToPrim $ 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 must be distinct: `" ++ show s ++ "`" @@ -240,7 +241,7 @@ VGM.write iter v $ i + 1 (!to, !iRevEdge, !_) <- ACIGV.read (gG VG.! v) i levelTo <- VGM.read level to - revCap <- readCapacity gG to iRevEdge + revCap <- readCapacityST gG to iRevEdge if levelV <= levelTo || revCap == 0 then loop res else do @@ -248,8 +249,8 @@ if d <= 0 then loop res -- no flow. ignore else do - modifyCapacity (gG VG.! v) (+ d) i - modifyCapacity (gG VG.! to) (subtract d) iRevEdge + modifyCapacityST (gG VG.! v) (+ d) i + modifyCapacityST (gG VG.! to) (subtract d) iRevEdge let !res' = res + d if res' == up then pure res' @@ -314,7 +315,7 @@ Int -> -- | Minimum cut m (VU.Vector Bit) -minCut MfGraph {..} s = do +minCut MfGraph {..} s = stToPrim $ do visited <- VUM.replicate nG $ Bit False que <- ACIQ.new nG -- we could use a growable queue here ACIQ.pushBack que s @@ -352,12 +353,12 @@ Int -> -- | Tuple of @(from, to, cap, flow)@ m (Int, Int, cap, cap) -getEdge MfGraph {..} i = do +getEdge MfGraph {..} i = stToPrim $ do m <- ACIGV.length posG let !_ = ACIA.checkEdge "AtCoder.MaxFlow.getEdge" i m (!from, !iEdge) <- ACIGV.read posG i (!to, !iRevEdge, !cap) <- ACIGV.read (gG VG.! from) iEdge - revCap <- readCapacity gG to iRevEdge + revCap <- readCapacityST gG to iRevEdge pure (from, to, cap + revCap, revCap) -- | Returns the current internal state of the edges: @(from, to, cap, flow)@. The edges are ordered @@ -400,32 +401,32 @@ -- | New flow cap -> m () -changeEdge MfGraph {..} i newCap newFlow = do +changeEdge MfGraph {..} i newCap newFlow = stToPrim $ do m <- ACIGV.length posG let !_ = ACIA.checkEdge "AtCoder.MaxFlow.changeEdge" i m let !_ = ACIA.runtimeAssert (0 <= newFlow && newFlow <= newCap) "AtCoder.MaxFlow.changeEdge: invalid flow or capacity" -- not Show (!from, !iEdge) <- ACIGV.read posG i (!to, !iRevEdge, !_) <- ACIGV.read (gG VG.! from) iEdge - writeCapacity gG from iEdge $! newCap - newFlow - writeCapacity gG to iRevEdge $! newFlow + writeCapacityST gG from iEdge $! newCap - newFlow + writeCapacityST gG to iRevEdge $! newFlow -- | \(O(1)\) Internal helper. -{-# INLINE readCapacity #-} -readCapacity :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => V.Vector (ACIGV.GrowVec (PrimState m) (Int, Int, cap)) -> Int -> Int -> m cap -readCapacity gvs v i = do +{-# INLINE readCapacityST #-} +readCapacityST :: (Num cap, Ord cap, VU.Unbox cap) => V.Vector (ACIGV.GrowVec s (Int, Int, cap)) -> Int -> Int -> ST s cap +readCapacityST gvs v i = do (VUM.MV_3 _ _ _ c) <- readMutVar $ ACIGV.vecGV $ gvs VG.! v VGM.read c i -- | \(O(1)\) Internal helper. -{-# INLINE writeCapacity #-} -writeCapacity :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => V.Vector (ACIGV.GrowVec (PrimState m) (Int, Int, cap)) -> Int -> Int -> cap -> m () -writeCapacity gvs v i cap = do +{-# INLINE writeCapacityST #-} +writeCapacityST :: (Num cap, Ord cap, VU.Unbox cap) => V.Vector (ACIGV.GrowVec s (Int, Int, cap)) -> Int -> Int -> cap -> ST s () +writeCapacityST gvs v i cap = do (VUM.MV_3 _ _ _ c) <- readMutVar $ ACIGV.vecGV $ gvs VG.! v VGM.write c i cap -- | \(O(1)\) Internal helper. -{-# INLINE modifyCapacity #-} -modifyCapacity :: (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) => ACIGV.GrowVec (PrimState m) (Int, Int, cap) -> (cap -> cap) -> Int -> m () -modifyCapacity gv f i = do +{-# INLINE modifyCapacityST #-} +modifyCapacityST :: (Num cap, Ord cap, VU.Unbox cap) => ACIGV.GrowVec s (Int, Int, cap) -> (cap -> cap) -> Int -> ST s () +modifyCapacityST gv f i = do (VUM.MV_3 _ _ _ c) <- readMutVar $ ACIGV.vecGV gv VUM.modify c f i diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -14,6 +14,7 @@ import Tests.Extra.MultiSet qualified import Tests.Extra.Semigroup.Matrix qualified import Tests.Extra.Semigroup.Permutation qualified +import Tests.Extra.Seq qualified import Tests.Extra.WaveletMatrix qualified import Tests.Extra.WaveletMatrix.BitVector qualified import Tests.Extra.WaveletMatrix.Raw qualified @@ -55,6 +56,7 @@ testGroup "MultiSet" Tests.Extra.MultiSet.tests, testGroup "Semigroup.Matrix" Tests.Extra.Semigroup.Matrix.tests, testGroup "Semigroup.Permutation" Tests.Extra.Semigroup.Permutation.tests, + testGroup "Seq" Tests.Extra.Seq.tests, testGroup "WaveletMatrix" Tests.Extra.WaveletMatrix.tests, testGroup "WaveletMatrix.BitVector" Tests.Extra.WaveletMatrix.BitVector.tests, testGroup "WaveletMatrix.Raw" Tests.Extra.WaveletMatrix.Raw.tests, diff --git a/test/Tests/Convolution.hs b/test/Tests/Convolution.hs --- a/test/Tests/Convolution.hs +++ b/test/Tests/Convolution.hs @@ -72,23 +72,27 @@ unit_butterfly :: TestTree unit_butterfly = testCase "butterfly" $ do - let modInt :: Int -> AM.ModInt998244353 - modInt = AM.new let expected = VU.fromList [10, 998244351, 173167434, 825076915] - vec <- VU.unsafeThaw $ VU.map modInt $ VU.fromList [1, 2, 3, 4] - info <- ACIC.newInfo @_ @998244353 - ACIC.butterfly info vec - (expected @=?) =<< VU.unsafeFreeze vec + let xs = VU.create $ do + let modInt :: Int -> AM.ModInt998244353 + modInt = AM.new + vec <- VU.unsafeThaw $ VU.map modInt $ VU.fromList [1, 2, 3, 4] + info <- ACIC.newInfo @_ @998244353 + ACIC.butterfly info vec + pure vec + expected @?= xs unit_invButterfly :: TestTree unit_invButterfly = testCase "invButterfly" $ do - let modInt :: Int -> AM.ModInt998244353 - modInt = AM.new let expected = VU.fromList [10, 911660634, 998244349, 86583717] - vec <- VU.unsafeThaw $ VU.map modInt $ VU.fromList [1, 2, 3, 4] - info <- ACIC.newInfo @_ @998244353 - ACIC.butterflyInv info vec - (expected @=?) =<< VU.unsafeFreeze vec + let xs = VU.create $ do + let modInt :: Int -> AM.ModInt998244353 + modInt = AM.new + vec <- VU.unsafeThaw $ VU.map modInt $ VU.fromList [1, 2, 3, 4] + info <- ACIC.newInfo @_ @998244353 + ACIC.butterflyInv info vec + pure vec + expected @=? xs testWithRangeMint :: forall p. diff --git a/test/Tests/Extra/IntervalMap.hs b/test/Tests/Extra/IntervalMap.hs --- a/test/Tests/Extra/IntervalMap.hs +++ b/test/Tests/Extra/IntervalMap.hs @@ -53,7 +53,7 @@ data Query = Contains Int - | Intersects (Int, Int) + | ContainsInterval (Int, Int) | Lookup (Int, Int) | -- | Read (Int, Int) ReadMaybe (Int, Int) @@ -80,7 +80,7 @@ queryGen n = do QC.oneof [ Contains <$> keyGen, - Intersects <$> intervalGen n, + ContainsInterval <$> intervalGen n, Lookup <$> intervalGen n, ReadMaybe <$> intervalGen n, Insert <$> intervalGen n <*> valGen, @@ -113,7 +113,7 @@ case q of Contains i -> do pure . B $ VU.any (\(!l, !r, !_) -> l <= i && i < r) intervals - Intersects (!l, !r) + ContainsInterval (!l, !r) | l >= r -> pure $ B False | otherwise -> pure . B $ VU.any (\(!l', !r', !_) -> l' <= l && r <= r') intervals Lookup (!l, !r) @@ -148,8 +148,8 @@ 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 + ContainsInterval (!l, !r) -> do + B <$> ITM.containsInterval itm l r Lookup (!l, !r) -> do LRX <$> ITM.lookup itm l r ReadMaybe (!l, !r) -> do diff --git a/test/Tests/Extra/Seq.hs b/test/Tests/Extra/Seq.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/Seq.hs @@ -0,0 +1,335 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.Seq (tests) where + +import AtCoder.Extra.Monoid.Affine1 (Affine1 (..)) +import AtCoder.Extra.Monoid.Affine1 qualified as Affine1 +import AtCoder.Extra.Pool qualified as P +import AtCoder.Extra.Seq qualified as Seq +import AtCoder.Internal.Assert qualified as ACIA +import Control.Monad (foldM_, when) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (RealWorld, runST) +import Data.Foldable (toList) +import Data.List qualified as L +import Data.Semigroup (Sum (..)) +import Data.Sequence qualified as S +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import System.IO.Unsafe (unsafePerformIO) +import Test.Hspec +import Test.QuickCheck.Monadic as QCM +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.Hspec +import Test.Tasty.QuickCheck as QC +import Tests.Util +import Prelude hiding (seq) + +unit_empty :: TestTree +unit_empty = testCase "empty sequence operations" $ do + let n = 4 -- TODO: randomize + seq <- Seq.new @_ @() @() n + h <- Seq.newNode seq () + h1 <- Seq.newSeq seq VU.empty + h2 <- Seq.newSeq seq VU.empty + h3 <- Seq.newSeq seq VU.empty + h4 <- Seq.newSeq seq VU.empty + + -- merge null and null + Seq.merge seq h1 h2 + assertBool "" . P.nullIndex =<< VGM.read (Seq.unHandle h1) 0 + assertBool "" . P.nullIndex =<< VGM.read (Seq.unHandle h2) 0 + + Seq.merge3 seq h1 h2 h3 + assertBool "" . P.nullIndex =<< VGM.read (Seq.unHandle h1) 0 + assertBool "" . P.nullIndex =<< VGM.read (Seq.unHandle h2) 0 + assertBool "" . P.nullIndex =<< VGM.read (Seq.unHandle h3) 0 + + Seq.merge4 seq h1 h2 h3 h4 + assertBool "" . P.nullIndex =<< VGM.read (Seq.unHandle h1) 0 + assertBool "" . P.nullIndex =<< VGM.read (Seq.unHandle h2) 0 + assertBool "" . P.nullIndex =<< VGM.read (Seq.unHandle h3) 0 + assertBool "" . P.nullIndex =<< VGM.read (Seq.unHandle h4) 0 + + -- merge a sequence and null + Seq.merge seq h1 h2 + assertBool "" . (== P.Index 0) =<< VGM.read (Seq.unHandle h) 0 + + Seq.merge3 seq h1 h2 h3 + assertBool "" . (== P.Index 0) =<< VGM.read (Seq.unHandle h) 0 + + Seq.merge4 seq h1 h2 h3 h4 + assertBool "" . (== P.Index 0) =<< VGM.read (Seq.unHandle h) 0 + +spec_boundaries :: IO TestTree +spec_boundaries = testSpec "boundaries" $ do + let n = 12 -- TODO: randomize + seq <- runIO $ Seq.new @_ @() @() n + h0 <- runIO $ Seq.newSeq seq VU.empty + + describe "split null at zero" $ do + r1 <- runIO $ Seq.split seq h0 0 + it "split - l" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle h0) 0 + it "split - r" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle r1) 0 + (!r2, !r3) <- runIO $ Seq.split3 seq h0 0 0 + it "split3 - l" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle h0) 0 + it "split3 - m" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle r2) 0 + it "split3 - r" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle r3) 0 + (!r4, !r5, !r6) <- runIO $ Seq.split4 seq h0 0 0 0 + it "split4 - a" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle h0) 0 + it "split4 - b" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle r4) 0 + it "split4 - c" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle r5) 0 + it "split4 - d" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle r6) 0 + + describe "split a sequence into a null and some" $ do + h <- runIO $ Seq.newSeq seq $ VU.replicate 1 () + i <- runIO $ VGM.read (Seq.unHandle h) 0 + r <- runIO $ Seq.split seq h 0 + it "a" $ (`shouldSatisfy` P.nullIndex) =<< VGM.read (Seq.unHandle h) 0 + it "b" $ (`shouldSatisfy` (== i)) =<< VGM.read (Seq.unHandle r) 0 + runIO $ Seq.free seq h + + let withSeq len f = do + h <- Seq.newSeq seq $ VU.replicate len () + _ <- f h + Seq.free seq h + pure () + + describe "split bounds (length 1)" $ do + it "l" $ withSeq 1 $ \h -> (`shouldThrow` anyException) $ Seq.split seq h (-1) + it "r" $ withSeq 1 $ \h -> (`shouldThrow` anyException) $ Seq.split seq h 2 + + -- TODO: assert that h1 == 0 + describe "split bounds (length 2)" $ do + it "l" $ withSeq 2 $ \h -> (`shouldThrow` anyException) $ Seq.split seq h (-1) + it "r" $ withSeq 2 $ \h -> (`shouldThrow` anyException) $ Seq.split seq h 3 + + pure () + +data Init = Init + { n :: {-# UNPACK #-} !Int, + q :: {-# UNPACK #-} !Int, + ref0 :: !(S.Seq (Sum Int)), + seqM :: !(IO (Seq.Seq RealWorld (Affine1 Int) (Sum Int), Seq.Handle RealWorld)) + } + +instance Show Init where + show Init {..} = show (n, ref0) + +instance QC.Arbitrary Init where + arbitrary = do + n <- QC.chooseInt (1, 64) + q <- QC.chooseInt (1, 5 * n) + pure $ Init n q (S.replicate n mempty) $ do + seq <- Seq.new (n + q) + root <- Seq.newSeq seq (VU.replicate n mempty) + pure (seq, root) + +data Query + = Reset + | Read !Int + | ReadMaybe !Int + | Write !Int !(Sum Int) + | Modify !Int !(Sum Int) + | Exchange !Int !(Sum Int) + | Prod !(Int, Int) + | ProdMaybe !(Int, Int) + | ProdAll + | ApplyIn !(Int, Int) !(Affine1 Int) + | ApplyToRoot !(Affine1 Int) + | -- | Reverse + Insert !Int !(Sum Int) + | Delete !Int + | Delete_ !Int + | -- | LowerBound (Sum Int) + LowerBoundProd !(Sum Int) + | Freeze + deriving (Show) + +-- | Arbitrary return type for the `Query` result. +data Result + = None + | B !Bool + | I !Int + | S !(Sum Int) + | MS !(Maybe (Sum Int)) + | F !(VU.Vector (Sum Int)) + deriving (Show, Eq) + +queryGen :: Int -> QC.Gen Query +queryGen n = do + QC.frequency + [ (rare, pure Reset), + (half, Read <$> keyGen), + (half, ReadMaybe <$> maybeKeyGen), + (often, Write <$> keyGen <*> valGen), + (often, Modify <$> keyGen <*> valGen), + (often, Exchange <$> keyGen <*> valGen), + (half, Prod <$> intervalGen n), + (half, ProdMaybe <$> maybeIntervalGen), + (often, pure ProdAll), + (often, ApplyIn <$> intervalGen n <*> fGen), + (often, ApplyToRoot <$> fGen), + (often, Insert <$> QC.chooseInt (0, n) <*> valGen), + (half, Delete <$> keyGen), + (half, Delete_ <$> keyGen), + (often, LowerBoundProd <$> valGen), + (rare, pure Freeze) + ] + where + rare = 1 + often = 10 + half = 5 + keyGen = QC.chooseInt (0, n - 1) + maybeKeyGen = QC.chooseInt (-1, n) + maybeIntervalGen = (,) <$> QC.chooseInt (-1, n + 1) <*> QC.chooseInt (-1, n + 1) + -- use non-negative values for monotoniously increasing sum + valGen = Sum <$> QC.chooseInt (0, 10) + -- NOTE: it might throw an error on overflow: + fGen = Affine1.new <$> QC.chooseInt (0, 4) <*> QC.chooseInt (0, 4) + +-- | containers. (referencial implementation) +handleRef :: S.Seq (Sum Int) -> Query -> (S.Seq (Sum Int), Result) +handleRef seq q = case q of + Reset -> (S.empty, None) + Read k -> (seq, S (S.index seq k)) + ReadMaybe k -> (seq, MS (seq S.!? k)) + Write k v -> (S.adjust (const v) k seq, None) + Modify k dx -> (S.adjust (<> dx) k seq, None) + Exchange k v -> (S.adjust (const v) k seq, S (S.index seq k)) + Prod (!l, !r) -> (seq, prod l r) + ProdMaybe (!l, !r) + | ACIA.testInterval l r n -> (seq, prodMaybe l r) + | otherwise -> (seq, MS Nothing) + ProdAll -> (seq, prod 0 (S.length seq)) + ApplyIn (!l, !r) f -> (apply l r f, None) + ApplyToRoot f -> (apply 0 (S.length seq) f, None) + Insert k v -> (S.insertAt k v seq, None) + Delete k -> (S.deleteAt k seq, S (S.index seq k)) + Delete_ k -> (S.deleteAt k seq, None) + LowerBoundProd x -> (seq, I (ilowerBound x)) + Freeze -> (seq, F (VU.fromList (toList seq))) + where + n = S.length seq + prod l r = S $ L.foldl' (<>) mempty $ map (S.index seq) [l .. r - 1] + prodMaybe l r = MS . Just $ L.foldl' (<>) mempty $ map (S.index seq) [l .. r - 1] + apply :: Int -> Int -> Affine1 Int -> S.Seq (Sum Int) + apply l r f = S.fromList . zipWith g [0 :: Int ..] $ toList seq + where + g i x + | l <= i && i < r = Sum $ Affine1.act f (getSum x) + | otherwise = x + ilowerBound x = length . takeWhile (<= x) . tail . L.scanl' (<>) mempty $ toList seq + +-- | ac-library-hs. +handleAcl :: (HasCallStack, PrimMonad m) => Seq.Seq (PrimState m) (Affine1 Int) (Sum Int) -> Seq.Handle (PrimState m) -> Query -> m Result +handleAcl seq handle q = case q of + Reset -> do + Seq.reset seq + Seq.invalidateHandle handle + pure None + Read k -> S <$> Seq.read seq handle k + ReadMaybe k -> MS <$> Seq.readMaybe seq handle k + Write k v -> do + Seq.write seq handle k v + pure None + Modify k dx -> do + Seq.modify seq handle (dx <>) k + pure None + Exchange k v -> do + S <$> Seq.exchange seq handle k v + Prod (!l, !r) -> do + S <$> Seq.prod seq handle l r + ProdMaybe (!l, !r) -> do + MS <$> Seq.prodMaybe seq handle l r + ProdAll -> do + S <$> Seq.prodAll seq handle + ApplyIn (!l, !r) f -> do + Seq.applyIn seq handle l r f + pure None + ApplyToRoot f -> do + Seq.applyToRoot seq handle f + pure None + Insert k v -> do + Seq.insert seq handle k v + pure None + Delete k -> do + S <$> Seq.delete seq handle k + Delete_ k -> do + Seq.delete_ seq handle k + pure None + LowerBoundProd x -> I <$> Seq.ilowerBoundProd seq handle (\_ y -> y <= x) + Freeze -> F <$> Seq.freeze seq handle + +-- | Ensures the capacity limit. +filterQuery :: S.Seq (Sum Int) -> Query -> Bool +filterQuery seq q = case q of + (Read k) -> idx k + (Write k _) -> idx k + (Modify k _) -> idx k + (Exchange k _) -> idx k + (Prod (!l, !r)) -> itv l r + (ApplyIn (!l, !r) _) -> itv l r + (Insert k _) -> 0 <= k && k <= n + (Delete k) -> idx k + (Delete_ k) -> idx k + _ -> True + where + n = S.length seq + idx k = 0 <= k && k < n + itv l r = 0 <= l && l <= r && r < n + +prop_randomTest :: Init -> QC.Property +prop_randomTest Init {..} = QCM.monadicIO $ do + (!seq, !root) <- QCM.run seqM + foldM_ + ( \ref _ -> do + query <- QCM.pick $ do + if S.null ref + then -- most operations throw an error for an empty sequence, so insert some value first: + Insert 0 . Sum <$> QC.chooseInt (0, 10) + else queryGen $ S.length ref + if filterQuery ref query + then do + -- run the query + let (!ref', !expected) = handleRef ref query + res <- QCM.run $ handleAcl seq root query + QCM.assertWith (expected == res) $ show (query, expected, res) + pure ref' + else pure ref + ) + ref0 + [0 .. q - 1] + +prop_bisectIndex :: QC.Gen QC.Property +prop_bisectIndex = do + n <- QC.chooseInt (1, 64) + k <- QC.chooseInt (0, n) + -- The higher order functinos for bisection method must take the index of intereseted node as a + -- argument + pure $ runST $ do + seq <- Seq.new n + root <- Seq.newSeq @_ @() seq $ VU.generate n Sum + lastRight1 <- VUM.replicate 1 (0 :: Int) + i1 <- Seq.ilowerBoundM seq root $ \i _ -> do + when (i < k) $ do + VGM.write lastRight1 0 $ i + 1 + pure $ i < k + lastRight2 <- VUM.replicate 1 (0 :: Int) + i2 <- Seq.ilowerBoundProdM seq root $ \i _ -> do + when (i < k) $ do + VGM.write lastRight2 0 $ i + 1 + pure $ i < k + i3 <- VGM.read lastRight1 0 + i4 <- VGM.read lastRight2 0 + pure . QC.conjoin $ map (QC.=== k) [i1, i2, i3, i4] + +tests :: [TestTree] +tests = + [ unit_empty, + unsafePerformIO spec_boundaries, + QC.testProperty "random test" prop_randomTest, + QC.testProperty "bisect index" prop_bisectIndex + ]