diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,28 +1,34 @@ # Revision history for acl-hs +## 1.2.2.0 -- Feb 2025 + +- Added `Extra.KdTree` and `Extra.LazyKdTree`. +- Added `clear` function to the dynamic segment tree family. +- Fixed `Extra.Hld.new` for a tree with a single vertex. + ## 1.2.1.0 -- Feb 2025 -- Added dynamic segment family -- Added `Extra.Seq.Map` -- Fixed `Extra.Pool.size` -- `Handle` is moved from `Extra.Seq` to `Extra.Pool` +- Added dynamic segment tree family. +- Added `Extra.Seq.Map`. +- Fixed `Extra.Pool.size`. +- `Handle` is moved from `Extra.Seq` to `Extra.Pool`. ## 1.2.0.0 -- Feb 2025 -- Added `AtCoder.Extra.Seq` -- Tweaked `INLINE` settings for less compile time +- 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` - - `SegAct` implementation for `Extra.Monoid.RangeAdd` over `Max` and `Min` were fixed + - `Matrix.diag` now does not take length parameter. + - `Extra.Math.primitiveRoot` is renamed to `primitiveRoot32`. + - `Internal.Convolution` functions now use `ST` instead of `PrimMonad`. + - `SegAct` implementation for `Extra.Monoid.RangeAdd` over `Max` and `Min` were fixed. ## 1.1.1.0 -- Jan 2025 -- Added `AtCoder.Extra.Tree.Lct` -- Added `blockCut`, `blockCutComponents` in `AtCoder.Extra.Graph` -- Added `popBack_` in `AtCoder.Internal.Buffer` -- Added `square`, `rank`, `inv`, `invRaw`, `detMod`, `detMint` in `AtCoder.Extra.Matrix` +- Added `AtCoder.Extra.Tree.Lct`. +- Added `blockCut`, `blockCutComponents` in `AtCoder.Extra.Graph`. +- Added `popBack_` in `AtCoder.Internal.Buffer`. +- Added `square`, `rank`, `inv`, `invRaw`, `detMod`, `detMint` in `AtCoder.Extra.Matrix`. ## 1.1.0.0 -- Jan 2025 diff --git a/ac-library-hs.cabal b/ac-library-hs.cabal --- a/ac-library-hs.cabal +++ b/ac-library-hs.cabal @@ -4,7 +4,7 @@ -- PVP summary: +-+------- breaking API changes -- | | +----- non-breaking API additions -- | | | +--- code changes with no API change -version: 1.2.1.0 +version: 1.2.2.0 synopsis: Data structures and algorithms description: Haskell port of [ac-library](https://github.com/atcoder/ac-library), a library for competitive @@ -43,18 +43,7 @@ , vector-algorithms <0.10 , wide-word <0.2 - default-language: GHC2021 - - -- compatible with GHC2024: - default-extensions: - DataKinds - DerivingStrategies - DisambiguateRecordFields - ExplicitNamespaces - GADTs - LambdaCase - MonoLocalBinds - RoleAnnotations + default-language: GHC2021 library import: warnings @@ -67,16 +56,18 @@ AtCoder.Extra.DynLazySegTree.Persistent AtCoder.Extra.DynLazySegTree.Raw AtCoder.Extra.DynSegTree - AtCoder.Extra.DynSegTree.Raw AtCoder.Extra.DynSegTree.Persistent + AtCoder.Extra.DynSegTree.Raw AtCoder.Extra.DynSparseSegTree - AtCoder.Extra.DynSparseSegTree.Raw AtCoder.Extra.DynSparseSegTree.Persistent + AtCoder.Extra.DynSparseSegTree.Raw AtCoder.Extra.Graph AtCoder.Extra.HashMap AtCoder.Extra.IntervalMap AtCoder.Extra.IntMap AtCoder.Extra.IntSet + AtCoder.Extra.KdTree + AtCoder.Extra.LazyKdTree AtCoder.Extra.Math AtCoder.Extra.Monoid AtCoder.Extra.Monoid.Affine1 @@ -154,6 +145,8 @@ Tests.Extra.IntervalMap Tests.Extra.IntMap Tests.Extra.IntSet + Tests.Extra.KdTree + Tests.Extra.LazyKdTree Tests.Extra.Math Tests.Extra.Monoid Tests.Extra.MultiSet diff --git a/benchmarks/BenchLib/ModInt/ModIntNats.hs b/benchmarks/BenchLib/ModInt/ModIntNats.hs --- a/benchmarks/BenchLib/ModInt/ModIntNats.hs +++ b/benchmarks/BenchLib/ModInt/ModIntNats.hs @@ -1,4 +1,5 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE TypeFamilies #-} diff --git a/benchmarks/BenchLib/ModInt/Modulus.hs b/benchmarks/BenchLib/ModInt/Modulus.hs --- a/benchmarks/BenchLib/ModInt/Modulus.hs +++ b/benchmarks/BenchLib/ModInt/Modulus.hs @@ -1,4 +1,5 @@ {-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE TypeFamilies #-} @@ -13,9 +14,9 @@ import BenchLib.MulMod.BarrettWideWord qualified as BarrettWideWord import Data.Bits import Data.Coerce (coerce) -import Data.Tagged (Tagged(..)) import Data.Proxy (Proxy) import Data.Ratio (denominator, numerator) +import Data.Tagged (Tagged (..)) import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Primitive qualified as P diff --git a/benchmarks/Tests/MulMod.hs b/benchmarks/Tests/MulMod.hs --- a/benchmarks/Tests/MulMod.hs +++ b/benchmarks/Tests/MulMod.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE DataKinds #-} + -- | A bit tedious tests. module Tests.MulMod (tests) where diff --git a/examples/LazySegTree.hs b/examples/LazySegTree.hs --- a/examples/LazySegTree.hs +++ b/examples/LazySegTree.hs @@ -1,6 +1,7 @@ -- This is a possible template of LazySegTree monoid instances for AtCoder contests. -- This code is copy-pasted to the lazy segtree document's example. +{-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE TypeFamilies #-} import AtCoder.LazySegTree qualified as LST diff --git a/src/AtCoder/Convolution.hs b/src/AtCoder/Convolution.hs --- a/src/AtCoder/Convolution.hs +++ b/src/AtCoder/Convolution.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DataKinds #-} {-# LANGUAGE MagicHash #-} -- | It calculates \((+,\times)\) convolution. Given two arrays \(a_0, a_1, \cdots, a_{N - 1}\) and \(b_0, b_1, \cdots, b_{M - 1}\), it calculates the array \(c\) of length \(N + M - 1\), defined by @@ -9,6 +10,7 @@ -- ==== __Example__ -- The convolution module basically works with `AtCoder.ModInt`: -- +-- >>> :set -XDataKinds -- >>> import AtCoder.Convolution qualified as C -- >>> import AtCoder.ModInt qualified as M -- >>> import Data.Proxy (Proxy) diff --git a/src/AtCoder/Extra/Bisect.hs b/src/AtCoder/Extra/Bisect.hs --- a/src/AtCoder/Extra/Bisect.hs +++ b/src/AtCoder/Extra/Bisect.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE LambdaCase #-} + -- | Bisection methods and binary search functions. They partition a half-open interval \([l, r)\) -- into two and return either the left or the right point of the boundary. -- diff --git a/src/AtCoder/Extra/DynLazySegTree.hs b/src/AtCoder/Extra/DynLazySegTree.hs --- a/src/AtCoder/Extra/DynLazySegTree.hs +++ b/src/AtCoder/Extra/DynLazySegTree.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DisambiguateRecordFields #-} {-# LANGUAGE TypeFamilies #-} -- | A dynamic, lazily propagated segment tree that covers a half-open interval \([l_0, r_0)\). @@ -78,6 +79,9 @@ maxRightM, -- -- * Conversions -- freeze, + + -- * Clear + clear, ) where @@ -307,3 +311,11 @@ maxRightM :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Raw.DynLazySegTree (PrimState m) f a -> P.Index -> (a -> m Bool) -> m Int maxRightM dst root f = do Raw.maxRightM dst root f + +-- | \(O(\log L)\) Claers all the nodes from the storage. +-- +-- @since 1.2.2.0 +{-# INLINE clear #-} +clear :: (PrimMonad m) => Raw.DynLazySegTree (PrimState m) f a -> m () +clear dst = do + P.clear (Raw.poolLdst dst) diff --git a/src/AtCoder/Extra/DynLazySegTree/Persistent.hs b/src/AtCoder/Extra/DynLazySegTree/Persistent.hs --- a/src/AtCoder/Extra/DynLazySegTree/Persistent.hs +++ b/src/AtCoder/Extra/DynLazySegTree/Persistent.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DisambiguateRecordFields #-} {-# LANGUAGE TypeFamilies #-} -- | A dynamic, persistent, lazily propagated segment tree that covers a half-open interval @@ -81,6 +82,9 @@ maxRightM, -- -- * Conversions -- freeze, + + -- * Clear + clear, ) where @@ -289,3 +293,11 @@ maxRightM :: (HasCallStack, PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Raw.DynLazySegTree (PrimState m) f a -> P.Index -> (a -> m Bool) -> m Int maxRightM dst root f = do Raw.maxRightM dst root f + +-- | \(O(\log L)\) Claers all the nodes from the storage. +-- +-- @since 1.2.2.0 +{-# INLINE clear #-} +clear :: (PrimMonad m) => Raw.DynLazySegTree (PrimState m) f a -> m () +clear dst = do + P.clear (Raw.poolLdst dst) diff --git a/src/AtCoder/Extra/DynLazySegTree/Raw.hs b/src/AtCoder/Extra/DynLazySegTree/Raw.hs --- a/src/AtCoder/Extra/DynLazySegTree/Raw.hs +++ b/src/AtCoder/Extra/DynLazySegTree/Raw.hs @@ -1,5 +1,6 @@ {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_HADDOCK hide #-} -- | Base module of a dynamic, lazily propagated segment tree. -- diff --git a/src/AtCoder/Extra/DynSegTree.hs b/src/AtCoder/Extra/DynSegTree.hs --- a/src/AtCoder/Extra/DynSegTree.hs +++ b/src/AtCoder/Extra/DynSegTree.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DisambiguateRecordFields #-} {-# LANGUAGE TypeFamilies #-} -- | A dynamic segment tree that covers a half-open interval \([l_0, r_0)\). Nodes are @@ -50,6 +51,7 @@ modifyM, -- exchange, -- read, + -- readMaybe, -- * Products prod, @@ -64,6 +66,9 @@ maxRightM, -- -- * Conversions -- freeze, + + -- * Clear + clear, ) where @@ -217,3 +222,11 @@ maxRightM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Raw.DynSegTree (PrimState m) a -> P.Index -> (a -> m Bool) -> m Int maxRightM dst root f = do Raw.maxRightM dst root f + +-- | \(O(\log L)\) Claers all the nodes from the storage. +-- +-- @since 1.2.2.0 +{-# INLINE clear #-} +clear :: (PrimMonad m) => Raw.DynSegTree (PrimState m) a -> m () +clear dst = do + P.clear (Raw.poolDst dst) diff --git a/src/AtCoder/Extra/DynSegTree/Persistent.hs b/src/AtCoder/Extra/DynSegTree/Persistent.hs --- a/src/AtCoder/Extra/DynSegTree/Persistent.hs +++ b/src/AtCoder/Extra/DynSegTree/Persistent.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DisambiguateRecordFields #-} {-# LANGUAGE TypeFamilies #-} -- | A dynamic, persistent segment tree that covers a half-open interval \([l_0, r_0)\). Nodes are @@ -63,6 +64,9 @@ maxRightM, -- -- * Conversions -- freeze, + + -- * Clear + clear, ) where @@ -210,3 +214,11 @@ maxRightM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Raw.DynSegTree (PrimState m) a -> P.Index -> (a -> m Bool) -> m Int maxRightM dst root f = do Raw.maxRightM dst root f + +-- | \(O(\log L)\) Claers all the nodes from the storage. +-- +-- @since 1.2.2.0 +{-# INLINE clear #-} +clear :: (PrimMonad m) => Raw.DynSegTree (PrimState m) a -> m () +clear dst = do + P.clear (Raw.poolDst dst) diff --git a/src/AtCoder/Extra/DynSegTree/Raw.hs b/src/AtCoder/Extra/DynSegTree/Raw.hs --- a/src/AtCoder/Extra/DynSegTree/Raw.hs +++ b/src/AtCoder/Extra/DynSegTree/Raw.hs @@ -1,5 +1,6 @@ {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_HADDOCK hide #-} -- | Base module of a dynamic segment tree. -- diff --git a/src/AtCoder/Extra/DynSparseSegTree.hs b/src/AtCoder/Extra/DynSparseSegTree.hs --- a/src/AtCoder/Extra/DynSparseSegTree.hs +++ b/src/AtCoder/Extra/DynSparseSegTree.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE DisambiguateRecordFields #-} +{-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE TypeFamilies #-} -- | A dynamic, sparse segment tree that covers a half-open interval \([l_0, r_0)\). Nodes are @@ -60,6 +62,9 @@ maxRightM, -- -- * Conversions -- freeze, + + -- * Clear + clear, ) where @@ -183,3 +188,11 @@ maxRightM dst (P.Handle handle) f = do root <- VGM.read handle 0 Raw.maxRightM dst root f + +-- | \(O(\log L)\) Claers all the nodes from the storage. +-- +-- @since 1.2.2.0 +{-# INLINE clear #-} +clear :: (PrimMonad m) => Raw.DynSparseSegTree (PrimState m) a -> m () +clear dst = do + P.clear (Raw.poolDsst dst) diff --git a/src/AtCoder/Extra/DynSparseSegTree/Persistent.hs b/src/AtCoder/Extra/DynSparseSegTree/Persistent.hs --- a/src/AtCoder/Extra/DynSparseSegTree/Persistent.hs +++ b/src/AtCoder/Extra/DynSparseSegTree/Persistent.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE DisambiguateRecordFields #-} +{-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE TypeFamilies #-} -- | A dynamic, sparse, persitent segment tree that covers a half-open interval \([l_0, r_0)\). Nodes are @@ -59,6 +61,9 @@ maxRightM, -- -- * Conversions -- freeze, + + -- * Clear + clear, ) where @@ -166,3 +171,11 @@ maxRightM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Raw.DynSparseSegTree (PrimState m) a -> P.Index -> (a -> m Bool) -> m Int maxRightM dst root f = do Raw.maxRightM dst root f + +-- | \(O(\log L)\) Claers all the nodes from the storage. +-- +-- @since 1.2.2.0 +{-# INLINE clear #-} +clear :: (PrimMonad m) => Raw.DynSparseSegTree (PrimState m) a -> m () +clear dst = do + P.clear (Raw.poolDsst dst) diff --git a/src/AtCoder/Extra/DynSparseSegTree/Raw.hs b/src/AtCoder/Extra/DynSparseSegTree/Raw.hs --- a/src/AtCoder/Extra/DynSparseSegTree/Raw.hs +++ b/src/AtCoder/Extra/DynSparseSegTree/Raw.hs @@ -1,5 +1,6 @@ {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_HADDOCK hide #-} -- | Base module of a dynamic, sparse segment tree. -- diff --git a/src/AtCoder/Extra/Graph.hs b/src/AtCoder/Extra/Graph.hs --- a/src/AtCoder/Extra/Graph.hs +++ b/src/AtCoder/Extra/Graph.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE LambdaCase #-} + -- | Re-export of the @Csr@ module and generic graph search functions. -- -- @since 1.1.0.0 @@ -58,7 +60,7 @@ -- -- @since 1.1.0.0 {-# INLINE swapDupe #-} -swapDupe :: (VU.Unbox (Int, Int, w)) => VU.Vector (Int, Int, w) -> VU.Vector (Int, Int, w) +swapDupe :: (VU.Unbox w) => VU.Vector (Int, Int, w) -> VU.Vector (Int, Int, w) swapDupe = VU.concatMap (\(!u, !v, !w) -> VU.fromListN 2 [(u, v, w), (v, u, w)]) -- | \(O(n)\) Converts non-directed edges into directional edges. This is a convenient function for @@ -86,7 +88,7 @@ -- -- @since 1.1.0.0 {-# INLINE swapDupe' #-} -swapDupe' :: (VU.Unbox (Int, Int)) => VU.Vector (Int, Int) -> VU.Vector (Int, Int) +swapDupe' :: VU.Vector (Int, Int) -> VU.Vector (Int, Int) swapDupe' = VU.concatMap (\(!u, !v) -> VU.fromListN 2 [(u, v), (v, u)]) -- | \(O(n + m)\) Returns the strongly connected components. diff --git a/src/AtCoder/Extra/IntMap.hs b/src/AtCoder/Extra/IntMap.hs --- a/src/AtCoder/Extra/IntMap.hs +++ b/src/AtCoder/Extra/IntMap.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE RecordWildCards #-} -- | A dense, fast `Int` map implemented as a 64-ary tree that covers an interval \([0, n)\). diff --git a/src/AtCoder/Extra/KdTree.hs b/src/AtCoder/Extra/KdTree.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/KdTree.hs @@ -0,0 +1,238 @@ +{-# LANGUAGE RecordWildCards #-} + +-- | Static, \(k\)-dimensional tree \((k = 2)\). +-- +-- - Points are fixed on `build`. +-- - Multiple points can exist at the same coordinate. +-- +-- ==== __Examples__ +-- >>> import AtCoder.Extra.KdTree qualified as KT +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let xys = VU.fromList [(0, 0), (1, 1), (4, 2)] +-- >>> let kt = KT.build2 xys +-- >>> -- Find point indices in [0, 2) x [0, 2) with maximum capacity 3 +-- >>> KT.findPointsIn kt 0 2 0 2 3 +-- [0,1] +-- +-- >>> KT.findNearestPoint kt 3 3 +-- Just 2 +-- +-- @since 1.2.2.0 +module AtCoder.Extra.KdTree + ( -- * K-dimensional tree + KdTree (..), + + -- * Constructors + build, + build2, + findPointsIn, + findNearestPoint, + ) +where + +import AtCoder.Internal.Assert qualified as ACIA +import AtCoder.Internal.Bit qualified as ACIB +import Control.Monad.ST (runST) +import Data.Bits +import Data.Ord (comparing) +import Data.Vector.Algorithms.Intro qualified as VAI +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) + +-- | Static, \(k\)-dimensional tree \((k = 2)\). +-- +-- @since 1.2.2.0 +data KdTree = KdTree + { -- | The number of points in the \(k\)-d tree. + -- + -- @since 1.2.2.0 + nKt :: {-# UNPACK #-} !Int, + -- | Rectangle information: inclusive (closed) ranges \([x_1, x_2) \times [y_1, y_2)\). + -- + -- @since 1.2.2.0 + incRectsKt :: !(VU.Vector (Int, Int, Int, Int)), + -- | Maps rectangle index to original point index. + -- + -- @since 1.2.2.0 + dataKt :: !(VU.Vector Int) + } + +-- | \(O(n \log n)\) Creates a `KdTree` from \(x\) and \(y\) vectors. +-- +-- ==== Constraints +-- - \(|\mathrm{xs}| = |\mathrm{ys}|\). +-- +-- @since 1.2.2.0 +{-# INLINEABLE build #-} +build :: + (HasCallStack) => + -- | \(x\) coordnates + VU.Vector Int -> + -- | \(y\) coordnates + VU.Vector Int -> + -- | `KdTree` + KdTree +build xs0 ys0 = + let nKt = VU.length xs0 + !_ = ACIA.runtimeAssert (nKt == VU.length ys0) "AtCoder.Extra.KdTree.buildST: the length of `xs`, `ys` and `vs` must be equal" + in if nKt == 0 + then KdTree 0 VU.empty VU.empty + else runST $ do + let vs0 = VU.generate nKt id + let logKt = countTrailingZeros $ ACIB.bitCeil (nKt + 1) + dat <- VUM.replicate (bit (logKt + 1)) (-1 :: Int) + incRectsVec <- VUM.replicate (bit (logKt + 1)) (maxBound, minBound, maxBound, minBound) + let VUM.MV_4 _ xMins xMaxes yMins yMaxes = incRectsVec + + -- - idx: rectangle index (one-based) + -- - xs, ys, vs: point information (x, y and monoid value) + -- - ids: maps sorted vertices to the original vertex indices + -- - divX: represents hyperplane direction for point partition + let -- buildSubtree :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int -> Bool -> ST s () + buildSubtree idx xs ys vs ids divX = do + let n = VU.length xs + + -- retrieve the bounds: + let (!xMin, !xMax, !yMin, !yMax) = + VU.foldl' + (\(!a, !b, !c, !d) (!x, !y) -> (min a x, max b x, min c y, max d y)) + (maxBound, minBound, maxBound, minBound) + $ VU.zip xs ys + VGM.modify xMins (min xMin) idx + VGM.modify xMaxes (max xMax) idx + VGM.modify yMins (min yMin) idx + VGM.modify yMaxes (max yMax) idx + + if n == 1 + then do + -- it's a terminal + VGM.write dat idx $ vs VG.! 0 + else do + -- partition the vertices into two: + let m = n `div` 2 + let is = VU.create $ do + vec <- VUM.generate n id + if divX + then VAI.selectBy (comparing (xs VG.!)) vec m + else VAI.selectBy (comparing (ys VG.!)) vec m + pure vec + + -- TODO: permute in-place? + let (!xsL, !xsR) = VG.splitAt m $ VG.backpermute xs is + let (!ysL, !ysR) = VG.splitAt m $ VG.backpermute ys is + let (!vsL, !vsR) = VG.splitAt m $ VG.backpermute vs is + let (!idsL, !idsR) = VG.splitAt m $ VG.backpermute ids is + + -- build the subtree: + buildSubtree (2 * idx + 0) xsL ysL vsL idsL (not divX) + buildSubtree (2 * idx + 1) xsR ysR vsR idsR (not divX) + + buildSubtree 1 xs0 ys0 vs0 (VU.generate nKt id) True + dataKt <- VU.unsafeFreeze dat + incRectsKt <- VU.unsafeFreeze incRectsVec + pure KdTree {..} + +-- | \(O(n \log n)\) Creates `KdTree` from a \((x, y)\) vector. +-- +-- ==== Constraints +-- - \(|\mathrm{xs}| = |\mathrm{ys}|\). +-- +-- @since 1.2.2.0 +{-# INLINE build2 #-} +build2 :: + (HasCallStack) => + -- | \(x, y\) coordnates + VU.Vector (Int, Int) -> + -- | `KdTree` + KdTree +build2 xys = build xs ys + where + (!xs, !ys) = VU.unzip xys + +-- | \(O(n \log n)\) Collects points in \([x_l, x_r) \times [y_l, y_r)\). +-- +-- @since 1.2.2.0 +{-# INLINE findPointsIn #-} +findPointsIn :: + (HasCallStack) => + -- | `KdTree` + KdTree -> + -- | \(x_l\) + Int -> + -- | \(x_r\) + Int -> + -- | \(y_l\) + Int -> + -- | \(y_r\) + Int -> + -- | Maximum number of points in \([x_l, x_r) \times [y_l, y_r)\). + Int -> + -- | Point indices in \([x_l, x_r) \times [y_l, y_r)\). + VU.Vector Int +findPointsIn KdTree {..} x1 x2 y1 y2 capacity + | nKt == 0 = VU.empty + | otherwise = runST $ do + res <- VUM.unsafeNew $ min nKt capacity + let inner i iPush + -- not intersected + | x2 <= xMin || xMax < x1 = pure iPush + | y2 <= yMin || yMax < y1 = pure iPush + -- a leaf + | vi /= -1 = do + VGM.write res iPush vi + pure $ iPush + 1 + -- a parental rectangle area + | otherwise = do + iPush' <- inner (2 * i + 0) iPush + inner (2 * i + 1) iPush' + where + (!xMin, !xMax, !yMin, !yMax) = incRectsKt VG.! i + vi = dataKt VG.! i + n <- inner 1 0 + VU.take n <$> VU.unsafeFreeze res + where + !_ = ACIA.runtimeAssert (x1 <= x2 && y1 <= y2) "AtCoder.Extra.KdTree.findPointsIn: given invalid interval" + +-- | \(O(\log n)\), only if the points are randomly distributed. Returns the index of the nearest +-- point, or `Nothing` if the `KdTree` has no point. +-- +-- @since 1.2.2.0 +{-# INLINE findNearestPoint #-} +findNearestPoint :: + (HasCallStack) => + -- | `KdTree` + KdTree -> + -- | \(x\) + Int -> + -- | \(y\) + Int -> + -- | The nearest point index + Maybe Int +findNearestPoint KdTree {..} x y + | nKt == 0 = Nothing + | otherwise = Just . fst $! inner 1 {- FIXME: -} (-1, -1) + where + clamp a aMin aMax = min aMax $ max a aMin + -- Used for pruning. It's |(x, y)|^2 if the (x, y) is within the rectangle. + bestDistSquared i = + let (!xMin, !xMax, !yMin, !yMax) = incRectsKt VG.! i + dx = x - clamp x xMin xMax + dy = y - clamp y yMin yMax + in dx * dx + dy * dy + -- returns (index, bestDist) + inner i res@(!resV, !resD) + -- pruning (we have a better point than any point in this rectangle) + | resV /= -1 && resD <= d = res + -- it's a leaf + | dataI /= -1 = (dataI, d) + -- look into the children + | d0 < d1 = inner (2 * i + 0) $ inner (2 * i + 1) res + | otherwise = inner (2 * i + 1) $ inner (2 * i + 0) res + where + d = bestDistSquared i + dataI = dataKt VG.! i + d0 = bestDistSquared (2 * i + 0) + d1 = bestDistSquared (2 * i + 0) diff --git a/src/AtCoder/Extra/LazyKdTree.hs b/src/AtCoder/Extra/LazyKdTree.hs new file mode 100644 --- /dev/null +++ b/src/AtCoder/Extra/LazyKdTree.hs @@ -0,0 +1,427 @@ +{-# LANGUAGE RecordWildCards #-} + +-- | Static, \(k\)-dimensional tree \((k = 2)\) with lazily propagated monoid actions and +-- commutative monoids. +-- +-- - Point coordinates are fixed on `build`. +-- - Multiple points can exist at the same coordinate. +-- +-- ==== __Examples__ +-- >>> import AtCoder.Extra.LazyKdTree qualified as LKT +-- >>> import AtCoder.Extra.Monoid.Affine1 (Affine1) +-- >>> import AtCoder.Extra.Monoid.Affine1 qualified as Affine1 +-- >>> import Data.Semigroup (Sum (..)) +-- >>> import Data.Vector.Unboxed qualified as VU +-- >>> let xyws = VU.fromList [(0, 0, Sum 1), (1, 1, Sum 2), (4, 2, Sum 3)] +-- >>> lkt <- LKT.build3 @_ @(Affine1 Int) @(Sum Int) xyws +-- +-- >>> -- Get monoid product in [0, 2) x [0, 2) +-- >>> LKT.prod lkt 0 2 0 2 +-- Sum {getSum = 3} +-- +-- >>> LKT.applyIn lkt 0 2 0 2 $ Affine1.new 2 1 +-- >>> LKT.prod lkt 0 2 0 2 +-- Sum {getSum = 8} +-- +-- >>> LKT.write lkt 0 $ Sum 10 +-- >>> LKT.prod lkt 0 2 0 2 +-- Sum {getSum = 15} +-- +-- @since 1.2.2.0 +module AtCoder.Extra.LazyKdTree + ( -- * K-dimensional tree + LazyKdTree (..), + + -- * Re-exports + SegAct (..), + + -- * Constructors + build, + build2, + build3, + + -- * Write + write, + modify, + modifyM, + + -- * Monoid products + prod, + allProd, + + -- * Apply + applyIn, + ) +where + +import AtCoder.Internal.Assert qualified as ACIA +import AtCoder.Internal.Bit qualified as ACIB +import AtCoder.LazySegTree (SegAct (..)) +import Control.Monad (unless, when) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) +import Control.Monad.ST (ST) +import Data.Bits +import Data.Foldable (for_) +import Data.Maybe (fromMaybe) +import Data.Ord (comparing) +import Data.Vector.Algorithms.Intro qualified as VAI +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) + +-- | Static, \(k\)-dimensional tree \((k = 2)\) with lazily propagated monoid actions and +-- commutative monoids. +-- +-- @since 1.2.2.0 +data LazyKdTree s f a = LazyKdTree + { -- | The number of points in the \(k\)-d tree. + -- + -- @since 1.2.2.0 + nLkt :: {-# UNPACK #-} !Int, + -- | \(\lceil \log_2 (n + 1) \rceil\) + -- + -- @since 1.2.2.0 + logLkt :: {-# UNPACK #-} !Int, + -- | Rectangle information: inclusive (closed) ranges \([x_1, x_2) \times [y_1, y_2)\). + -- + -- @since 1.2.2.0 + incRectsLkt :: !(VU.Vector (Int, Int, Int, Int)), + -- | Rectangle information: monoid values. + -- + -- @since 1.2.2.0 + dataLkt :: !(VUM.MVector s a), + -- | Rectangle information: lazily propagated monoid actions for children. + -- + -- @since 1.2.2.0 + lazyLkt :: !(VUM.MVector s f), + -- | Rectangle information: the number of vertices in the rectangle. + -- + -- @since 1.2.2.0 + sizeLkt :: !(VU.Vector Int), + -- | Maps original vertices into the belonging rectangle index. + -- + -- @since 1.2.2.0 + posLkt :: !(VU.Vector Int) + } + +-- | \(O(n \log n)\) Creates a `LazyKdTree` from @xs@, @ys@ and @ws@ vectors. +-- +-- ==== Constraints +-- - \(|\mathrm{xs}| = |\mathrm{ys}| = |\mathrm{vs}|\). +-- +-- @since 1.2.2.0 +{-# INLINE build #-} +build :: + (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) => + -- | \(x\) coordnates + VU.Vector Int -> + -- | \(y\) coordnates + VU.Vector Int -> + -- | monoid \(v\)alues + VU.Vector a -> + -- | `LazyKdTree` + m (LazyKdTree (PrimState m) f a) +build xs ys vs = stToPrim $ buildST xs ys vs + +-- | \(O(n \log n)\) Creates a `LazyKdTree` from @xys@ and @ws@ vectors. +-- +-- ==== Constraints +-- - \(|\mathrm{xys}| = |\mathrm{vs}|\). +-- +-- @since 1.2.2.0 +{-# INLINE build2 #-} +build2 :: + (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) => + -- | \((x, y)\) coordinates + VU.Vector (Int, Int) -> + -- | Monoid \(v\)alues + VU.Vector a -> + -- | `LazyKdTree` + m (LazyKdTree (PrimState m) f a) +build2 xys ws = stToPrim $ buildST xs ys ws + where + (!xs, !ys) = VU.unzip xys + +-- | \(O(n \log n)\) Creates a `LazyKdTree` from a @xyws@ vector. +-- +-- @since 1.2.2.0 +{-# INLINE build3 #-} +build3 :: + (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) => + -- | \((x, y, v)\) tuples + VU.Vector (Int, Int, a) -> + -- | `LazyKdTree` + m (LazyKdTree (PrimState m) f a) +build3 xyws = stToPrim $ buildST xs ys ws + where + (!xs, !ys, !ws) = VU.unzip3 xyws + +-- | \(O(\log n)\) Writes to the \(k\)-th point's monoid value. +-- +-- @since 1.2.2.0 +{-# INLINE write #-} +write :: + (HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) => + -- | `LazyKdTree` + LazyKdTree (PrimState m) f a -> + -- | Original vertex index. + Int -> + -- | Monoid value + a -> + -- | Monadic tuple + m () +write kt i x = stToPrim $ modifyM kt (pure . const x) i + +-- | \(O(\log n)\) Modifies the \(k\)-th point's monoid value. +-- +-- @since 1.2.2.0 +{-# INLINE modify #-} +modify :: + (HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) => + -- | `LazyKdTree` + LazyKdTree (PrimState m) f a -> + -- | Creates a new monoid value from the old one. + (a -> a) -> + -- | Original vertex index. + Int -> + -- | Monadic tuple + m () +modify kt f i = stToPrim $ modifyM kt (pure . f) i + +-- | \(O(\log n)\) Modifies the \(k\)-th point's monoid value. +-- +-- @since 1.2.2.0 +{-# INLINEABLE modifyM #-} +modifyM :: + (HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) => + -- | `LazyKdTree` + LazyKdTree (PrimState m) f a -> + -- | Creates a new monoid value from the old one. + (a -> m a) -> + -- | Original vertex index. + Int -> + -- | Monadic tuple + m () +modifyM kt@LazyKdTree {..} f i0 = do + let i_ = posLkt VG.! i0 + -- propagate lazily propagated monoid actions from the root: + stToPrim $ for_ [logLkt, logLkt - 1 .. 1] $ \k -> do + pushST kt (i_ .>>. k) + VGM.modifyM dataLkt f i_ + -- update parents: + let inner i + | i <= 1 = pure () + | otherwise = do + let i' = i `div` 2 + xl <- VGM.read dataLkt (2 * i' + 0) + xr <- VGM.read dataLkt (2 * i' + 1) + VGM.write dataLkt i' $! xl <> xr + inner i' + stToPrim $ inner i_ + +-- | \(O(\log n)\) Returns monoid product in \([x_l, x_r) \times [y_l, y_r)\). +-- +-- @since 1.2.2.0 +{-# INLINE prod #-} +prod :: + (HasCallStack, PrimMonad m, Eq f, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | `LazyKdTree` + LazyKdTree (PrimState m) f a -> + -- | \(x_l\) + Int -> + -- | \(x_r\) + Int -> + -- | \(y_l\) + Int -> + -- | \(y_r\) + Int -> + -- | Monoid product in \([x_l, x_r) \times [y_l, y_r)\) + m a +prod kt x1 x2 y1 y2 = stToPrim $ prodST kt x1 x2 y1 y2 + +-- | \(O(1)\) Returns monoid product of all the points. +-- +-- @since 1.2.2.0 +{-# INLINE allProd #-} +allProd :: + (PrimMonad m, Monoid a, VU.Unbox a) => + -- | `LazyKdTree` + LazyKdTree (PrimState m) f a -> + -- | Monoid product in the whole space. + m a +allProd kt = do + -- In case of zero vertices, use `Maybe`: + fromMaybe mempty <$> VGM.readMaybe (dataLkt kt) 1 + +-- | \(O(\log n)\) Applies a monoid action to points in \([x_l, x_r) \times [y_l, y_r)\). +-- +-- @since 1.2.2.0 +{-# INLINE applyIn #-} +applyIn :: + (HasCallStack, PrimMonad m, Eq f, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => + -- | `LazyKdTree` + LazyKdTree (PrimState m) f a -> + -- | \(x_l\) + Int -> + -- | \(x_r\) + Int -> + -- | \(y_l\) + Int -> + -- | \(y_r\) + Int -> + -- | \(f\) + f -> + -- | Monadic tuple + m () +applyIn kt x1 x2 y1 y2 f = stToPrim $ applyInST kt 1 x1 x2 y1 y2 f + +-- ------------------------------------------------------------------------------------------------- +-- Private +-- ------------------------------------------------------------------------------------------------- + +{-# INLINEABLE buildST #-} +buildST :: forall s f a. (HasCallStack, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) => VU.Vector Int -> VU.Vector Int -> VU.Vector a -> ST s (LazyKdTree s f a) +buildST xs0 ys0 vs0 = do + let nLkt = VU.length xs0 + let !_ = ACIA.runtimeAssert (nLkt == VU.length ys0 && nLkt == VU.length vs0) "AtCoder.Extra.LazyKdTree.buildST: the length of `xs`, `ys` and `vs` must be equal" + if nLkt == 0 + then do + let logLkt = 0 + dataLkt <- VUM.new 0 + lazyLkt <- VUM.new 0 + let incRectsLkt = VU.empty + let sizeLkt = VU.empty + let posLkt = VU.empty + pure LazyKdTree {..} + else do + let logLkt = countTrailingZeros $ ACIB.bitCeil (nLkt + 1) + dataLkt <- VUM.unsafeNew (bit (logLkt + 1)) + lazyLkt <- VUM.replicate (bit logLkt) mempty + incRectsVec <- VUM.replicate (bit (logLkt + 1)) (maxBound, minBound, maxBound, minBound) + size <- VUM.unsafeNew (bit (logLkt + 1)) + pos <- VUM.unsafeNew nLkt + let VUM.MV_4 _ xMins xMaxes yMins yMaxes = incRectsVec + + -- - idx: rectangle index (one-based) + -- - xs, ys, vs: point information (x, y and monoid value) + -- - ids: maps sorted vertices to the original vertex indices + -- - divX: represents hyperplane direction for point partition + let buildSubtree :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector a -> VU.Vector Int -> Bool -> ST s () + buildSubtree idx xs ys vs ids divX = do + let n = VU.length xs + VGM.write size idx n + + -- retrieve the bounds: + let (!xMin, !xMax, !yMin, !yMax) = + VU.foldl' + (\(!a, !b, !c, !d) (!x, !y) -> (min a x, max b x, min c y, max d y)) + (maxBound, minBound, maxBound, minBound) + $ VU.zip xs ys + VGM.modify xMins (min xMin) idx + VGM.modify xMaxes (max xMax) idx + VGM.modify yMins (min yMin) idx + VGM.modify yMaxes (max yMax) idx + + if n == 1 + then do + -- it's a terminal. note that it's not always a leaf; the case is handled carefully in + -- other methods + VGM.write dataLkt idx $ vs VG.! 0 + -- record original vertex index -> rectangle index + VGM.write pos (ids VG.! 0) idx + else do + -- partition the vertices into two: + let m = n `div` 2 + let is = VU.create $ do + vec <- VUM.generate n id + if divX + then VAI.selectBy (comparing (xs VG.!)) vec m + else VAI.selectBy (comparing (ys VG.!)) vec m + pure vec + + -- TODO: permute in-place? + let (!xsL, !xsR) = VG.splitAt m $ VG.backpermute xs is + let (!ysL, !ysR) = VG.splitAt m $ VG.backpermute ys is + let (!vsL, !vsR) = VG.splitAt m $ VG.backpermute vs is + let (!idsL, !idsR) = VG.splitAt m $ VG.backpermute ids is + + -- build the subtree: + buildSubtree (2 * idx + 0) xsL ysL vsL idsL (not divX) + buildSubtree (2 * idx + 1) xsR ysR vsR idsR (not divX) + xl <- VGM.read dataLkt (2 * idx + 0) + xr <- VGM.read dataLkt (2 * idx + 1) + VGM.write dataLkt idx $! xl <> xr + + buildSubtree 1 xs0 ys0 vs0 (VU.generate nLkt id) True + sizeLkt <- VU.unsafeFreeze size + posLkt <- VU.unsafeFreeze pos + incRectsLkt <- VU.unsafeFreeze incRectsVec + pure LazyKdTree {..} + +{-# INLINE applyAtST #-} +applyAtST :: (SegAct f a, VU.Unbox f, VU.Unbox a) => LazyKdTree s f a -> Int -> f -> ST s () +applyAtST LazyKdTree {..} i f = do + -- NOTE: Here we're asssuming each monoid value has length one. If you need a monoid of length + -- zero, e.g., if you're just reserving new point insertion, you must not rely on + -- `segActWithLength`. You might want to use `V2` instead of `Sum`. + let len = sizeLkt VG.! i + VGM.modify dataLkt (segActWithLength len f) i + when (i < bit logLkt) $ do + VGM.modify lazyLkt (f <>) i + +-- TODO: consider `INLINE`? +{-# INLINE pushST #-} +pushST :: (SegAct f a, Eq f, VU.Unbox f, VU.Unbox a) => LazyKdTree s f a -> Int -> ST s () +pushST kt@LazyKdTree {..} i = do + lazy <- VGM.read lazyLkt i + unless (lazy == mempty) $ do + applyAtST kt (2 * i + 0) lazy + applyAtST kt (2 * i + 1) lazy + VGM.write lazyLkt i mempty + +{-# INLINEABLE prodST #-} +prodST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => LazyKdTree s f a -> Int -> Int -> Int -> Int -> ST s a +prodST kt@LazyKdTree {..} x1 x2 y1 y2 + | x1 >= x2 || y1 >= y2 = pure mempty + | otherwise = inner 1 + where + inner i = case incRectsLkt VG.!? i of + Nothing -> pure mempty + Just (!xl, !xr, !yl, !yr) + -- TODO: what is this? + | xl > xr -> pure mempty + -- not intersecting + | x2 <= xl || x1 > xr || y2 <= yl || y1 > yr -> pure mempty + -- the rectangle is fully contained by the query: + | x1 <= xl && xr < x2 && y1 <= yl && yr < y2 -> do + VGM.read dataLkt i + | otherwise -> do + pushST kt i + l <- inner (2 * i + 0) + r <- inner (2 * i + 1) + pure $! l <> r + +{-# INLINEABLE applyInST #-} +applyInST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => LazyKdTree s f a -> Int -> Int -> Int -> Int -> Int -> f -> ST s () +applyInST kt@LazyKdTree {..} i0 x1 x2 y1 y2 f + | x1 >= x2 || y1 >= y2 = pure () + | otherwise = inner i0 + where + inner i = case incRectsLkt VG.!? i of + Nothing -> pure mempty + Just (!xl, !xr, !yl, !yr) + | xl > xr -> pure () + -- not intersecting + | x2 <= xl || x1 > xr || y2 <= yl || y1 > yr -> pure () + -- the rectangle is fully contained by the query: + | x1 <= xl && xr < x2 && y1 <= yl && yr < y2 -> do + applyAtST kt i f + | otherwise -> do + pushST kt i + inner (2 * i + 0) + inner (2 * i + 1) + l <- VGM.read dataLkt (2 * i + 0) + r <- VGM.read dataLkt (2 * i + 1) + VGM.write dataLkt i $! l <> r diff --git a/src/AtCoder/Extra/Monoid/Affine1.hs b/src/AtCoder/Extra/Monoid/Affine1.hs --- a/src/AtCoder/Extra/Monoid/Affine1.hs +++ b/src/AtCoder/Extra/Monoid/Affine1.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE TypeFamilies #-} -- | Monoid action \(f: x \rightarrow ax + b\). @@ -23,6 +24,7 @@ where import AtCoder.Extra.Math qualified as ACEM +import AtCoder.Extra.Monoid.V2 qualified as V2 import AtCoder.LazySegTree (SegAct (..)) import Data.Coerce (coerce) import Data.Foldable (foldl') @@ -149,7 +151,15 @@ {-# INLINE segActWithLength #-} segActWithLength !len (Dual f) (Sum !x) = Sum $! actWithLength len (coerce f) x --- not works as SegAct for Product, Min, and Max. +-- | @since 1.2.2.0 +instance (Num a) => SegAct (Affine1 a) (V2.V2 a) where + {-# INLINE segAct #-} + segAct (Affine1 (!a, !b)) (V2.V2 (!v, !len)) = V2.V2 (a * v + b * len, len) + +-- | @since 1.2.2.0 +instance (Num a) => SegAct (Dual (Affine1 a)) (V2.V2 a) where + {-# INLINE segAct #-} + segAct (Dual (Affine1 (!a, !b))) (V2.V2 (!v, !len)) = V2.V2 (a * v + b * len, len) -- | @since 1.0.0.0 newtype instance VU.MVector s (Affine1 a) = MV_Affine1 (VU.MVector s (Affine1Repr a)) diff --git a/src/AtCoder/Extra/Monoid/Mat2x2.hs b/src/AtCoder/Extra/Monoid/Mat2x2.hs --- a/src/AtCoder/Extra/Monoid/Mat2x2.hs +++ b/src/AtCoder/Extra/Monoid/Mat2x2.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE TypeFamilies #-} -- | Monoid action \(f: x \rightarrow ax + b\). Less efficient than @Affine1@, but compatible with diff --git a/src/AtCoder/Extra/Monoid/RangeAdd.hs b/src/AtCoder/Extra/Monoid/RangeAdd.hs --- a/src/AtCoder/Extra/Monoid/RangeAdd.hs +++ b/src/AtCoder/Extra/Monoid/RangeAdd.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE TypeFamilies #-} -- | Monoid action \(f: x \rightarrow x + d\). diff --git a/src/AtCoder/Extra/Monoid/RangeSet.hs b/src/AtCoder/Extra/Monoid/RangeSet.hs --- a/src/AtCoder/Extra/Monoid/RangeSet.hs +++ b/src/AtCoder/Extra/Monoid/RangeSet.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE TypeFamilies #-} -- | Monoid action \(f: x \rightarrow a\). diff --git a/src/AtCoder/Extra/Monoid/V2.hs b/src/AtCoder/Extra/Monoid/V2.hs --- a/src/AtCoder/Extra/Monoid/V2.hs +++ b/src/AtCoder/Extra/Monoid/V2.hs @@ -1,16 +1,22 @@ +{-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE TypeFamilies #-} --- | A monoid acted on by `Mat2x2`, an affine transformation target. +-- | A monoid acted on by `Mat2x2` or `Affine1`, an affine transformation target. -- +-- ==== As an `Affine1` action target +-- Compared to `Sum`, `V2` hold the length in the second value. +-- -- @since 1.1.0.0 module AtCoder.Extra.Monoid.V2 ( -- * V2 V2 (..), V2Repr, - -- * Constructor + -- * Constructors new, unV2, + zero, + isZero, ) where @@ -38,19 +44,34 @@ -- @since 1.1.0.0 type V2Repr a = (a, a) --- | \(O(1)\) Creates `V2` of length \(1\). +-- | \(O(1)\) Creates a `V2` of length \(1\). -- -- @since 1.1.0.0 {-# INLINE new #-} new :: (Num a) => a -> V2 a new !a = V2 (a, 1) +-- | \(O(1)\) Creates a `V2` of length \(0\), i.e., `mempty`. Note that this value does not change +-- on affine transformation. +-- +-- @since 1.2.2.0 +{-# INLINE zero #-} +zero :: (Num a) => V2 a +zero = V2 (0, 0) + -- | \(O(1)\) Retrieves the value of `V2`, discarding the length information. -- -- @since 1.1.0.0 {-# INLINE unV2 #-} unV2 :: V2 a -> a unV2 (V2 (!a, !_)) = a + +-- | \(O(1)\) Returns whether the `V2` is equal to `zero` +-- +-- @since 1.2.2.0 +{-# INLINE isZero #-} +isZero :: (Num a, Eq a) => V2 a -> Bool +isZero = (== zero) -- | @since 1.1.0.0 instance (Num a) => Semigroup (V2 a) where 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 @@ -1,3 +1,4 @@ +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE RecordWildCards #-} -- | A fast, mutable multiset for `Int` keys backed by a @HashMap@. Most operations are performed diff --git a/src/AtCoder/Extra/Semigroup/Permutation.hs b/src/AtCoder/Extra/Semigroup/Permutation.hs --- a/src/AtCoder/Extra/Semigroup/Permutation.hs +++ b/src/AtCoder/Extra/Semigroup/Permutation.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE DerivingStrategies #-} + -- | A permutation represented by a vector, mainly for binary exponentiation. -- -- The permutation is a left semigroup action: \(p_2 (p_1 x) = (p_2 \circ p_1) x\). diff --git a/src/AtCoder/Extra/Seq/Map.hs b/src/AtCoder/Extra/Seq/Map.hs --- a/src/AtCoder/Extra/Seq/Map.hs +++ b/src/AtCoder/Extra/Seq/Map.hs @@ -77,6 +77,7 @@ -- ** Products prodInInterval, + -- TODO: prodIn -- ** Applications applyInInterval, diff --git a/src/AtCoder/Extra/Seq/Raw.hs b/src/AtCoder/Extra/Seq/Raw.hs --- a/src/AtCoder/Extra/Seq/Raw.hs +++ b/src/AtCoder/Extra/Seq/Raw.hs @@ -1,6 +1,7 @@ {-# LANGUAGE DerivingVia #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_HADDOCK hide #-} -- | Base module for implementing dynamic sequences. It internaly uses a splay tree and user has to -- track the root node change. @@ -84,6 +85,7 @@ freezeST, -- * Internals + -- | These functions are exported primarily for @Map@ implementations. splitMaxRightWithST, maxRightWithST, @@ -435,7 +437,7 @@ -- [ ) -- * root' (splayed) -- * rootL (detached from the root) - -- * rootL' (splayed) + -- \* rootL' (splayed) -- * right(rootL'): node that corresponds to [l, r) root' <- splayKthST seq root r rootL <- VGM.read lSeq $ coerce root' diff --git a/src/AtCoder/Extra/Tree/Hld.hs b/src/AtCoder/Extra/Tree/Hld.hs --- a/src/AtCoder/Extra/Tree/Hld.hs +++ b/src/AtCoder/Extra/Tree/Hld.hs @@ -75,6 +75,7 @@ -- >>> -- + -- >>> -- +--4--5 -- >>> let n = 6 +-- >>> -- note that the edges must be bi-directed: -- >>> let tree = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (2, 3), (1, 4), (4, 5)] -- >>> let hld = Hld.new tree -- @@ -215,107 +216,128 @@ -- | \(O(n)\) Creates an `Hld` with \(0\) as the root vertex. -- +-- ==== Constraints +-- - \(n \ge 1\) +-- - The input graph must be a tree; the edges must be non-directed (both \((u, v, w)\) and +-- \((v, u, w)\) edges are required). +-- -- @since 1.1.0.0 -{-# INLINABLE new #-} +{-# INLINEABLE new #-} new :: forall w. (HasCallStack) => Gr.Csr w -> Hld new tree = newAt tree 0 -- | \(O(n)\) Creates an `Hld` with a root vertex specified. -- +-- ==== Constraints +-- - \(n \ge 1\) +-- - The input graph must be a tree; the edges must be non-directed (both \((u, v, w)\) and +-- \((v, u, w)\) edges are required). +-- -- @since 1.1.0.0 -{-# INLINABLE newAt #-} +{-# INLINEABLE newAt #-} newAt :: forall w. (HasCallStack) => Gr.Csr w -> Vertex -> Hld -newAt tree root = runST $ do - -- Re-create adjacent vertices so that the biggest subtree's head vertex comes first. - -- - -- We /could/ instead record the biggest adjacent subtree vertex for each vertex, but the other - -- DFS would be harder. - let (!tree', !parent, !depths, !subtreeSize) = runST $ do - adjVec <- VU.thaw (Gr.adjCsr tree) - parent_ <- VUM.unsafeNew n - depths_ <- VUM.unsafeNew n - subtreeSize_ <- VUM.unsafeNew n +newAt tree root + | n == 1 = + Hld + 0 + (VU.singleton (-1)) + (VU.singleton 0) + (VU.singleton (-1)) + (VU.singleton 0) + (VU.singleton 0) + (VU.singleton 1) + | otherwise = runST $ do + -- Re-create adjacent vertices so that the biggest subtree's head vertex comes first. + -- + -- We /could/ instead record the biggest adjacent subtree vertex for each vertex, but the other + -- DFS would be harder. + let (!tree', !parent, !depths, !subtreeSize) = runST $ do + adjVec <- VU.thaw (Gr.adjCsr tree) + parent_ <- VUM.unsafeNew n + depths_ <- VUM.unsafeNew n + subtreeSize_ <- VUM.unsafeNew n - _ <- (\f -> fix f 0 (-1) root) $ \loop depth p v1 -> do - VGM.write parent_ v1 p - VGM.write depths_ v1 depth + _ <- (\f -> fix f 0 (-1) root) $ \loop depth p v1 -> do + VGM.write parent_ v1 p + VGM.write depths_ v1 depth - (!size1, !eBig) <- - VU.foldM' - ( \(!size1, !eBig) (!e2, !v2) -> do - if v2 == p - then pure (size1, eBig) - else do - size2 <- loop (depth + 1) v1 v2 - -- NOTE: It's `>` because we should swap at least once if there's some vertex other - -- that the parent_. - pure (size1 + size2, if size1 > size2 then eBig else e2) - ) - (1 :: Int, -1) - (tree `Gr.eAdj` v1) + (!size1, !eBig) <- + VU.foldM' + ( \(!size1, !eBig) (!e2, !v2) -> do + if v2 == p + then pure (size1, eBig) + else do + size2 <- loop (depth + 1) v1 v2 + -- NOTE: It's `>` because we should swap at least once if there's some vertex other + -- that the parent_. + pure (size1 + size2, if size1 > size2 then eBig else e2) + ) + (1 :: Int, -1) + (tree `Gr.eAdj` v1) - -- move the biggest subtree's head to the first adjacent vertex. - -- it means the "heavy edge" or the longest segment. - when (eBig /= -1) $ do - VGM.swap adjVec eBig $ fst (VG.head (tree `Gr.eAdj` v1)) + -- move the biggest subtree's head to the first adjacent vertex. + -- it means the "heavy edge" or the longest segment. + when (eBig /= -1) $ do + VGM.swap adjVec eBig $ fst (VG.head (tree `Gr.eAdj` v1)) - -- record subtree size - VGM.write subtreeSize_ v1 size1 + -- record subtree size + VGM.write subtreeSize_ v1 size1 - pure size1 + pure size1 - !vec <- VU.unsafeFreeze adjVec - (tree {Gr.adjCsr = vec},,,) - <$> VU.unsafeFreeze parent_ - <*> VU.unsafeFreeze depths_ - <*> VU.unsafeFreeze subtreeSize_ + !vec <- VU.unsafeFreeze adjVec + (tree {Gr.adjCsr = vec},,,) + <$> VU.unsafeFreeze parent_ + <*> VU.unsafeFreeze depths_ + <*> VU.unsafeFreeze subtreeSize_ - -- vertex -> reindexed vertex index - indices <- VUM.replicate n (-1 :: Int) + -- vertex -> reindexed vertex index + indices <- VUM.replicate n (-1 :: Int) - -- vertex -> head vertex of the segment - heads <- VUM.replicate n (-1 :: Int) + -- vertex -> head vertex of the segment + heads <- VUM.replicate n (-1 :: Int) - _ <- (\f -> fix f (0 :: Int) root (-1) root) $ \loop acc h p v1 -> do - -- reindex: - VGM.write indices v1 acc - let !acc' = acc + 1 + _ <- (\f -> fix f (0 :: Int) root (-1) root) $ \loop acc h p v1 -> do + -- reindex: + VGM.write indices v1 acc + let !acc' = acc + 1 - VGM.write heads v1 h + VGM.write heads v1 h - -- when the first vertex is within the same segment: - let (!adj1, !rest) = fromJust $ VU.uncons (tree' `Gr.adj` v1) - acc'' <- - if adj1 == p - then pure acc' - else loop acc' h v1 adj1 + -- when the first vertex is within the same segment: + let (!adj1, !rest) = fromJust $ VU.uncons (tree' `Gr.adj` v1) + acc'' <- + if adj1 == p + then pure acc' + else loop acc' h v1 adj1 - -- the others are in other segments: - VU.foldM' - ( \a v2 -> do - if v2 == p - then pure a - else loop a v2 v1 v2 - ) - acc'' - rest + -- the others are in other segments: + VU.foldM' + ( \a v2 -> do + if v2 == p + then pure a + else loop a v2 v1 v2 + ) + acc'' + rest - !indices' <- VU.unsafeFreeze indices - let !revIndex = VU.update (VU.replicate n (-1)) $ VU.imap (flip (,)) indices' + !indices' <- VU.unsafeFreeze indices + let !revIndex = VU.update (VU.replicate n (-1)) $ VU.imap (flip (,)) indices' - Hld root parent indices' - <$> VU.unsafeFreeze heads - <*> pure revIndex - <*> pure depths - <*> pure subtreeSize + Hld root parent indices' + <$> VU.unsafeFreeze heads + <*> pure revIndex + <*> pure depths + <*> pure subtreeSize where !n = Gr.nCsr tree !_ = ACIA.runtimeAssert (2 * (Gr.nCsr tree - 1) == Gr.mCsr tree) "AtCoder.Extra.Hld.newAt: not a non-directed tree" + !_ = ACIA.runtimeAssert (n >= 1) "AtCoder.Extra.Hld.newAt: the tree must have at least one vertex" -- | \(O(\log n)\) Calculates the lowest common ancestor of \(u\) and \(v\). -- -- @since 1.1.0.0 -{-# INLINABLE lca #-} +{-# INLINEABLE lca #-} lca :: (HasCallStack) => Hld -> Vertex -> Vertex -> Vertex lca Hld {..} = inner where @@ -338,7 +360,7 @@ -- is bigger than the depth of \(v\). -- -- @since 1.1.0.0 -{-# INLINABLE ancestor #-} +{-# INLINEABLE ancestor #-} ancestor :: (HasCallStack) => Hld -> Vertex -> Int -> Vertex ancestor Hld {..} parent k0 = inner parent k0 where @@ -357,7 +379,7 @@ -- Throws an error if `k` is out -- -- @since 1.1.0.0 -{-# INLINABLE jump #-} +{-# INLINEABLE jump #-} jump :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int -> Maybe Vertex jump hld@Hld {..} u v k | k > lenU + lenV = Nothing @@ -373,7 +395,7 @@ -- | \(O(\log n)\) Returns the length of the path between \(u\) and \(v\). -- -- @since 1.1.0.0 -{-# INLINABLE lengthBetween #-} +{-# INLINEABLE lengthBetween #-} lengthBetween :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int lengthBetween hld@Hld {..} u v = du - dLca + dv - dLca where @@ -385,7 +407,7 @@ -- | \(O(n)\) Returns the vertices on the path between \(u\) and \(v\). -- -- @since 1.1.0.0 -{-# INLINABLE path #-} +{-# INLINEABLE path #-} path :: (HasCallStack) => Hld -> Vertex -> Vertex -> [Vertex] path hld@Hld {..} u v = concatMap expand $ pathSegmentsInclusive WeightsAreOnVertices hld u v where @@ -400,7 +422,7 @@ -- `WeightsAreOnEdges`. This is the trick to put edge weights to on vertices. -- -- @since 1.1.0.0 -{-# INLINABLE pathSegmentsInclusive #-} +{-# INLINEABLE pathSegmentsInclusive #-} pathSegmentsInclusive :: (HasCallStack) => WeightPolicy -> Hld -> Vertex -> Vertex -> [(VertexHld, VertexHld)] pathSegmentsInclusive weightPolicy Hld {..} x0 y0 = done $ inner x0 [] y0 [] where @@ -440,7 +462,7 @@ -- corresponds to the subtree segments rooted at the given @subtreeRoot@. -- -- @since 1.1.0.0 -{-# INLINABLE subtreeSegmentInclusive #-} +{-# INLINEABLE subtreeSegmentInclusive #-} subtreeSegmentInclusive :: (HasCallStack) => Hld -> Vertex -> (VertexHld, VertexHld) subtreeSegmentInclusive Hld {..} subtreeRoot = (ir, ir + sr - 1) where @@ -450,7 +472,7 @@ -- | \(O(1)\) Returns `True` if \(u\) is in a subtree of \(r\). -- -- @since 1.1.0.0 -{-# INLINABLE isInSubtree #-} +{-# INLINEABLE isInSubtree #-} isInSubtree :: (HasCallStack) => Hld -> Vertex -> Vertex -> Bool isInSubtree hld@Hld {..} r_ u = l <= iu && iu <= r where diff --git a/src/AtCoder/Extra/Tree/TreeMonoid.hs b/src/AtCoder/Extra/Tree/TreeMonoid.hs --- a/src/AtCoder/Extra/Tree/TreeMonoid.hs +++ b/src/AtCoder/Extra/Tree/TreeMonoid.hs @@ -38,6 +38,7 @@ -- >>> -- + -- >>> -- +--4--5 -- >>> let n = 6 +-- >>> -- note that the edges must be bi-directed: -- >>> let tree = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (2, 3), (1, 4), (4, 5)] -- >>> let weights = VU.generate n Sum -- vertex `i` is given weight of `i` -- >>> let hld = Hld.new tree @@ -63,8 +64,11 @@ -- >>> -- +--4--5 -- >>> let n = 6 -- >>> let edges = VU.fromList [(0, 1, Sum (1 :: Int)), (1, 2, Sum 2), (2, 3, Sum 3), (1, 4, Sum 4), (4, 5, Sum 5)] +-- >>> -- note that the edges must be bi-directed: -- >>> let tree = Gr.build n $ Gr.swapDupe edges -- >>> let hld = Hld.new tree +-- >>> let hld = Hld.new tree +-- >>> -- note that the edge doesn't have to be bi-directed: -- >>> tm <- TM.fromEdges hld {- `Sum` is commutative -} Commute edges -- >>> TM.prod tm 1 3 -- Sum {getSum = 5} 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 @@ -1,4 +1,5 @@ {-# LANGUAGE RecordWildCards #-} +{-# OPTIONS_HADDOCK hide #-} -- original implementation: -- diff --git a/src/AtCoder/FenwickTree.hs b/src/AtCoder/FenwickTree.hs --- a/src/AtCoder/FenwickTree.hs +++ b/src/AtCoder/FenwickTree.hs @@ -43,16 +43,22 @@ -- * Adding add, - -- * Accessor + -- * Accessors sum, sumMaybe, + + -- * Bisection methods + maxRight, + maxRightM, + minLeft, + minLeftM, ) where import AtCoder.Internal.Assert qualified as ACIA import Control.Monad (when) import Control.Monad.Fix (fix) -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) import Data.Bits import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU @@ -170,3 +176,186 @@ xr <- prefixSum ft r xl <- prefixSum ft l pure $! xr - xl + +-- | (Extra API) Applies a binary search on the Fenwick tree. It returns an index \(r\) that +-- satisfies both of the following. +-- +-- - \(r = l\) or \(f(a[l] + a[l + 1] + ... + a[r - 1])\) returns `True`. +-- - \(r = n\) or \(f(a[l] + a[l + 1] + ... + a[r]))\) returns `False`. +-- +-- If \(f\) is monotone, this is the maximum \(r\) that satisfies +-- \(f(a[l] + a[l + 1] + ... + a[r - 1])\). +-- +-- ==== Constraints +-- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side effect. +-- - \(f(0)\) returns `True` +-- - \(0 \leq l \leq n\) +-- +-- ==== Complexity +-- - \(O(\log n)\) +-- +-- @since 1.2.2.0 +{-# INLINE maxRight #-} +maxRight :: + (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => + -- | The Fenwick tree + FenwickTree (PrimState m) a -> + -- | \(l\) + Int -> + -- | \(p\): user predicate + (a -> Bool) -> + -- | \(r\): \(p\) holds for \([l, r)\) + m Int +maxRight ft l0 f = maxRightM ft l0 (pure . f) + +-- | (Extra API) Monadic variant of `maxRight`. +-- +-- ==== Constraints +-- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side effect. +-- - \(f(0)\) returns `True` +-- - \(0 \leq l \leq n\) +-- +-- ==== Complexity +-- - \(O(\log n)\) +-- +-- @since 1.2.2.0 +{-# INLINEABLE maxRightM #-} +maxRightM :: + forall m a. + (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => + -- | The Fenwick tree + FenwickTree (PrimState m) a -> + -- | \(l\) + Int -> + -- | \(p\): user predicate + (a -> m Bool) -> + -- | \(r\): \(p\) holds for \([l, r)\) + m Int +maxRightM FenwickTree {..} l0 f = do + b0 <- f 0 + let !_ = ACIA.runtimeAssert b0 "AtCoder.FenwickTree.maxRightM: `f 0` must return `True`" + + let inner i !s + | odd i = do + ds <- stToPrim $ VGM.read dataFt (i - 1) + inner (i - 1) $! s - ds + | i == 0 = pure (i, 64 - countLeadingZeros nFt, s) + | i + bit k > nFt = pure (i, k, s) + | otherwise = do + t <- stToPrim $ (s +) <$> VGM.read dataFt (i + bit k - 1) + b <- f t + if not b + then pure (i, k, s) + else do + di <- stToPrim $ VGM.read dataFt (i - 1) + inner (i - i .&. (-i)) $! s - di + where + k = countTrailingZeros i - 1 + + -- we could start from an arbitrary l, but the API is limited to one + (!i0, !k0, !s0) <- inner l0 (0 :: a) + + let inner2 i k_ !s + | k_ == 0 = pure i + | i + bit k - 1 < VGM.length dataFt = do + t <- stToPrim $ (s +) <$> VGM.read dataFt (i + bit k - 1) + b <- f t + if b + then inner2 (i + bit k) k t + else inner2 i k s + | otherwise = inner2 i k s + where + k = k_ - 1 + + inner2 i0 k0 s0 + +-- | Applies a binary search on the Fenwick tree. It returns an index \(l\) that satisfies both of +-- the following. +-- +-- - \(l = r\) or \(f(a[l] + a[l + 1] + ... + a[r - 1])\) returns `True`. +-- - \(l = 0\) or \(f(a[l - 1] + a[l] + ... + a[r - 1])\) returns `False`. +-- +-- If \(f\) is monotone, this is the minimum \(l\) that satisfies +-- \(f(a[l] + a[l + 1] + ... + a[r - 1])\). +-- +-- ==== Constraints +-- +-- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side +-- effect. +-- - \(f(0)\) returns `True` +-- - \(0 \leq r \leq n\) +-- +-- ==== Complexity +-- - \(O(\log n)\) +-- +-- @since 1.2.2.0 +{-# INLINE minLeft #-} +minLeft :: + (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => + -- | The Fenwick tree + FenwickTree (PrimState m) a -> + -- | \(r\) + Int -> + -- | \(p\): user prediate + (a -> Bool) -> + -- | \(l\): \(p\) holds for \([l, r)\) + m Int +minLeft ft r0 f = minLeftM ft r0 (pure . f) + +-- | (Extra API) Monadic variant of `minLeft`. +-- +-- ==== Constraints +-- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side effect. +-- - \(f(0)\) returns `True` +-- - \(0 \leq l \leq n\) +-- +-- ==== Complexity +-- - \(O(\log n)\) +-- +-- @since 1.2.2.0 +{-# INLINEABLE minLeftM #-} +minLeftM :: + forall m a. + (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => + -- | The Fenwick tree + FenwickTree (PrimState m) a -> + -- | \(r\) + Int -> + -- | \(p\): user prediate + (a -> m Bool) -> + -- | \(l\): \(p\) holds for \([l, r)\) + m Int +minLeftM FenwickTree {..} r0 f = do + b0 <- f 0 + let !_ = ACIA.runtimeAssert b0 "AtCoder.FenwickTree.minLeftM: `f 0` must return `True`" + + let inner i k !s + | i <= 0 = pure (i, k, s) + | otherwise = do + b <- f s + if not b + then pure (i, k, s) + else do + s' <- stToPrim $ (s +) <$> VGM.read dataFt (i - 1) + inner (i - i .&. (-i)) (countTrailingZeros i) s' + + -- we could start from an arbitrary r, but the API is limited to n + (!i0, !k0, !s0) <- inner r0 (0 :: Int) (0 :: a) + + b0_ <- f s0 + if b0_ + then do + let !_ = ACIA.runtimeAssert (i0 == 0) "AtCoder.FenwickTree.minLeftM: implementation error" + pure 0 + else do + let inner2 i k_ !s + | k_ == 0 = pure i + | otherwise = do + let k = k_ - 1 + t <- stToPrim $ (s -) <$> VGM.read dataFt (i + bit k - 1) + b <- f t + if b + then inner2 i k s + else inner2 (i + bit k) k t + + (+ 1) <$> inner2 i0 k0 s0 diff --git a/src/AtCoder/Internal/Assert.hs b/src/AtCoder/Internal/Assert.hs --- a/src/AtCoder/Internal/Assert.hs +++ b/src/AtCoder/Internal/Assert.hs @@ -62,7 +62,7 @@ checkCustom, errorCustom, - -- * Interval assertion + -- * Interval assertions checkInterval, errorInterval, checkIntervalBounded, diff --git a/src/AtCoder/LazySegTree.hs b/src/AtCoder/LazySegTree.hs --- a/src/AtCoder/LazySegTree.hs +++ b/src/AtCoder/LazySegTree.hs @@ -143,7 +143,7 @@ -- -- If you implement `SegAct` via `segActWithLength`, satisfy one more propety: -- --- [Linear left monoid action] @'segActWithLength' len f a = 'Data.Semigroup.stimes' len ('segAct' f a) a@. +-- [Linear left monoid action] @'segActWithLength' len f (stimes len a) = 'Data.Semigroup.stimes' len ('segAct' f a)@. -- -- ==== Invariant -- In `SegAct` instances, new semigroup values are always given from the left: @new '<>' old@. The @@ -198,6 +198,7 @@ -- Define your monoid action @F@ and your acted monoid @X@: -- -- @ +-- {-# LANGUAGE DerivingStrategies #-} -- {-# LANGUAGE TypeFamilies #-} -- -- import AtCoder.LazySegTree qualified as LST @@ -732,11 +733,11 @@ -- - \(O(\log n)\) -- -- @since 1.0.0.0 -{-# INLINE maxRightM #-} +{-# INLINEABLE maxRightM #-} maxRightM :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> (a -> m Bool) -> m Int maxRightM self@LazySegTree {..} l0 g = do b <- g mempty - let !_ = ACIA.runtimeAssert b "AtCoder.LazySegTree.maxRightM: `g mempty` returned `False`" + let !_ = ACIA.runtimeAssert b "AtCoder.LazySegTree.maxRightM: `g mempty` must return `False`" if l0 == nLst then pure nLst else do @@ -749,9 +750,9 @@ !_ = ACIA.runtimeAssert (0 <= l0 && l0 <= nLst) $ "AtCoder.LazySegTree.maxRightM: given invalid `left` index `" ++ show l0 ++ "` over length `" ++ show nLst ++ "`" inner l !sm = do let l' = chooseBit l - !sm' <- (sm <>) <$> VGM.read dLst l' + !sm' <- stToPrim $ (sm <>) <$> VGM.read dLst l' b <- g sm' - if not $ b + if not b then do inner2 l' sm else do diff --git a/src/AtCoder/ModInt.hs b/src/AtCoder/ModInt.hs --- a/src/AtCoder/ModInt.hs +++ b/src/AtCoder/ModInt.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE TypeFamilies #-} @@ -158,6 +160,7 @@ -- | Retrieves the `Int` value from a `KnownNat`. -- +-- >>> :set -XDataKinds -- >>> import Data.Proxy (Proxy(..)) -- >>> modVal (Proxy @42) -- 42 @@ -169,6 +172,7 @@ -- | Retrieves the `Int` value from a `KnownNat`. -- +-- >>> :set -XDataKinds -- >>> :set -XMagicHash -- >>> import GHC.Exts (proxy#) -- >>> modVal# (proxy# @42) diff --git a/src/AtCoder/SegTree.hs b/src/AtCoder/SegTree.hs --- a/src/AtCoder/SegTree.hs +++ b/src/AtCoder/SegTree.hs @@ -105,7 +105,7 @@ import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Bit qualified as ACIBIT -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) import Data.Bits (countTrailingZeros, testBit, (.&.), (.>>.)) import Data.Foldable (for_) import Data.Vector.Generic.Mutable qualified as VGM @@ -361,7 +361,7 @@ -- - \(O(\log n)\) -- -- @since 1.0.0.0 -{-# INLINE minLeftM #-} +{-# INLINEABLE minLeftM #-} minLeftM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => -- | The segment tree @@ -383,7 +383,7 @@ !_ = ACIA.runtimeAssert (0 <= r0 && r0 <= nSt) $ "AtCoder.SegTree.minLeftM: given invalid `right` index `" ++ show r0 ++ "` over length `" ++ show nSt ++ "`" inner r !sm = do let r' = chooseBit $ r - 1 - !sm' <- (<> sm) <$> VGM.read dSt r' + !sm' <- stToPrim $ (<> sm) <$> VGM.read dSt r' b <- f sm' if not b then do @@ -398,7 +398,7 @@ inner2 r sm | r < sizeSt = do let r' = 2 * r + 1 - !sm' <- (<> sm) <$> VGM.read dSt r' + !sm' <- stToPrim $ (<> sm) <$> VGM.read dSt r' b <- f sm' if b then inner2 (r' - 1) sm' @@ -447,7 +447,7 @@ -- - \(O(\log n)\) -- -- @since 1.0.0.0 -{-# INLINE maxRightM #-} +{-# INLINEABLE maxRightM #-} maxRightM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => -- | The segment tree @@ -460,7 +460,7 @@ m Int maxRightM SegTree {..} l0 f = do b <- f mempty - let !_ = ACIA.runtimeAssert b "AtCoder.SegTree.maxRightM: `f mempty` returned `False`" + let !_ = ACIA.runtimeAssert b "AtCoder.SegTree.maxRightM: `f mempty` must return `True`" if l0 == nSt then pure nSt else inner (l0 + sizeSt) mempty @@ -469,7 +469,7 @@ !_ = ACIA.runtimeAssert (0 <= l0 && l0 <= nSt) $ "AtCoder.SegTree.maxRightM: given invalid `left` index `" ++ show l0 ++ "` over length `" ++ show nSt ++ "`" inner l !sm = do let l' = chooseBit l - !sm' <- (sm <>) <$> VGM.read dSt l' + !sm' <- stToPrim $ (sm <>) <$> VGM.read dSt l' b <- f sm' if not b then do @@ -486,7 +486,7 @@ inner2 l !sm | l < sizeSt = do let l' = 2 * l - !sm' <- (sm <>) <$> VGM.read dSt l' + !sm' <- stToPrim $ (sm <>) <$> VGM.read dSt l' b <- f sm' if b then inner2 (l' + 1) sm' diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -15,6 +15,8 @@ import Tests.Extra.IntMap qualified import Tests.Extra.IntSet qualified import Tests.Extra.IntervalMap qualified +import Tests.Extra.KdTree qualified +import Tests.Extra.LazyKdTree qualified import Tests.Extra.Math qualified import Tests.Extra.Monoid qualified import Tests.Extra.MultiSet qualified @@ -64,6 +66,8 @@ testGroup "IntervalMap" Tests.Extra.IntervalMap.tests, testGroup "IntMap" Tests.Extra.IntMap.tests, testGroup "IntSet" Tests.Extra.IntSet.tests, + testGroup "KdTree" Tests.Extra.KdTree.tests, + testGroup "LazyKdTree" Tests.Extra.LazyKdTree.tests, testGroup "Math" Tests.Extra.Math.tests, testGroup "Monoid" Tests.Extra.Monoid.tests, testGroup "MultiSet" Tests.Extra.MultiSet.tests, diff --git a/test/Tests/Convolution.hs b/test/Tests/Convolution.hs --- a/test/Tests/Convolution.hs +++ b/test/Tests/Convolution.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE DataKinds #-} + module Tests.Convolution (tests) where import AtCoder.Convolution qualified as ACC diff --git a/test/Tests/Extra/DynLazySegTree.hs b/test/Tests/Extra/DynLazySegTree.hs --- a/test/Tests/Extra/DynLazySegTree.hs +++ b/test/Tests/Extra/DynLazySegTree.hs @@ -11,7 +11,6 @@ import Control.Monad.ST (RealWorld, runST) import Data.Foldable (for_) import Data.Semigroup (Sum (..)) -import Data.Vector.Algorithms.Intro qualified as VAI import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM @@ -166,16 +165,15 @@ -- prop_foldl is tested with large array verification problem -prop_maxRight :: Int -> [QC.NonNegative Int] -> QC.Property -prop_maxRight xRef xs_ = - not (null xs_) QC.==> do - let xs = VU.modify VAI.sort $ VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ - expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs - res = runST $ do - seg <- Seg.new @_ @() @(Sum Int) (2 * VU.length xs) 0 (VU.length xs) - root <- Seg.newSeq seg $ VU.map Sum xs - Seg.maxRight seg root (<= Sum xRef) - in expected QC.=== res +prop_maxRight :: Int -> QC.NonEmptyList (QC.NonNegative Int) -> QC.Property +prop_maxRight xRef (QC.NonEmpty xs_) = + let xs = VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ + expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs + res = runST $ do + seg <- Seg.new @_ @() @(Sum Int) (2 * VU.length xs) 0 (VU.length xs) + root <- Seg.newSeq seg $ VU.map Sum xs + Seg.maxRight seg root (<= Sum xRef) + in expected QC.=== res tests :: [TestTree] tests = diff --git a/test/Tests/Extra/DynLazySegTree/Persistent.hs b/test/Tests/Extra/DynLazySegTree/Persistent.hs --- a/test/Tests/Extra/DynLazySegTree/Persistent.hs +++ b/test/Tests/Extra/DynLazySegTree/Persistent.hs @@ -12,7 +12,6 @@ import Control.Monad.ST (RealWorld, runST) import Data.Foldable (for_) import Data.Semigroup (Sum (..)) -import Data.Vector.Algorithms.Intro qualified as VAI import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM @@ -171,16 +170,15 @@ -- prop_foldl is tested with large array verification problem -prop_maxRight :: Int -> [QC.NonNegative Int] -> QC.Property -prop_maxRight xRef xs_ = - not (null xs_) QC.==> do - let xs = VU.modify VAI.sort $ VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ - expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs - res = runST $ do - seg <- Seg.new @_ @() @(Sum Int) (2 * VU.length xs) 0 (VU.length xs) - root <- Seg.newSeq seg $ VU.map Sum xs - Seg.maxRight seg root (<= Sum xRef) - in expected QC.=== res +prop_maxRight :: Int -> QC.NonEmptyList (QC.NonNegative Int) -> QC.Property +prop_maxRight xRef (QC.NonEmpty xs_) = + let xs = VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ + expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs + res = runST $ do + seg <- Seg.new @_ @() @(Sum Int) (2 * VU.length xs) 0 (VU.length xs) + root <- Seg.newSeq seg $ VU.map Sum xs + Seg.maxRight seg root (<= Sum xRef) + in expected QC.=== res tests :: [TestTree] tests = diff --git a/test/Tests/Extra/DynSegTree.hs b/test/Tests/Extra/DynSegTree.hs --- a/test/Tests/Extra/DynSegTree.hs +++ b/test/Tests/Extra/DynSegTree.hs @@ -8,7 +8,6 @@ import Control.Monad.ST (RealWorld, runST) import Data.Foldable (for_) import Data.Semigroup (Sum (..)) -import Data.Vector.Algorithms.Intro qualified as VAI import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM @@ -136,16 +135,15 @@ -- prop_foldl is tested with large array verification problem -prop_maxRight :: Int -> [QC.NonNegative Int] -> QC.Property -prop_maxRight xRef xs_ = - not (null xs_) QC.==> do - let xs = VU.modify VAI.sort $ VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ - expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs - res = runST $ do - seg <- Seg.new @_ @(Sum Int) (2 * VU.length xs) 0 (VU.length xs) - root <- Seg.newSeq seg $ VU.map Sum xs - Seg.maxRight seg root (<= Sum xRef) - in expected QC.=== res +prop_maxRight :: Int -> QC.NonEmptyList (QC.NonNegative Int) -> QC.Property +prop_maxRight xRef (QC.NonEmpty xs_) = + let xs = VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ + expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs + res = runST $ do + seg <- Seg.new @_ @(Sum Int) (2 * VU.length xs) 0 (VU.length xs) + root <- Seg.newSeq seg $ VU.map Sum xs + Seg.maxRight seg root (<= Sum xRef) + in expected QC.=== res tests :: [TestTree] tests = diff --git a/test/Tests/Extra/DynSegTree/Persistent.hs b/test/Tests/Extra/DynSegTree/Persistent.hs --- a/test/Tests/Extra/DynSegTree/Persistent.hs +++ b/test/Tests/Extra/DynSegTree/Persistent.hs @@ -9,7 +9,6 @@ import Control.Monad.ST (RealWorld, runST) import Data.Foldable (for_) import Data.Semigroup (Sum (..)) -import Data.Vector.Algorithms.Intro qualified as VAI import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM @@ -141,16 +140,15 @@ -- prop_foldl is tested with large array verification problem -prop_maxRight :: Int -> [QC.NonNegative Int] -> QC.Property -prop_maxRight xRef xs_ = - not (null xs_) QC.==> do - let xs = VU.modify VAI.sort $ VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ - expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs - res = runST $ do - seg <- Seg.new @_ @(Sum Int) (2 * VU.length xs) 0 (VU.length xs) - root <- Seg.newSeq seg $ VU.map Sum xs - Seg.maxRight seg root (<= Sum xRef) - in expected QC.=== res +prop_maxRight :: Int -> QC.NonEmptyList (QC.NonNegative Int) -> QC.Property +prop_maxRight xRef (QC.NonEmpty xs_) = + let xs = VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ + expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs + res = runST $ do + seg <- Seg.new @_ @(Sum Int) (2 * VU.length xs) 0 (VU.length xs) + root <- Seg.newSeq seg $ VU.map Sum xs + Seg.maxRight seg root (<= Sum xRef) + in expected QC.=== res tests :: [TestTree] tests = diff --git a/test/Tests/Extra/DynSparseSegTree.hs b/test/Tests/Extra/DynSparseSegTree.hs --- a/test/Tests/Extra/DynSparseSegTree.hs +++ b/test/Tests/Extra/DynSparseSegTree.hs @@ -8,7 +8,6 @@ import Control.Monad.ST (RealWorld, runST) import Data.Foldable (for_) import Data.Semigroup (Sum (..)) -import Data.Vector.Algorithms.Intro qualified as VAI import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM @@ -121,10 +120,10 @@ -- prop_foldl is tested with large array verification problem -prop_maxRight :: Int -> [QC.NonNegative Int] -> QC.Property -prop_maxRight xRef xs_ = +prop_maxRight :: Int -> QC.NonEmptyList (QC.NonNegative Int) -> QC.Property +prop_maxRight xRef (QC.NonEmpty xs_) = not (null xs_) QC.==> do - let xs = VU.modify VAI.sort $ VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ + let xs = VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs res = runST $ do seg <- Seg.new @_ @(Sum Int) (2 * VU.length xs) 0 (VU.length xs) diff --git a/test/Tests/Extra/DynSparseSegTree/Persistent.hs b/test/Tests/Extra/DynSparseSegTree/Persistent.hs --- a/test/Tests/Extra/DynSparseSegTree/Persistent.hs +++ b/test/Tests/Extra/DynSparseSegTree/Persistent.hs @@ -8,7 +8,6 @@ import Control.Monad.ST (RealWorld, runST) import Data.Foldable (for_) import Data.Semigroup (Sum (..)) -import Data.Vector.Algorithms.Intro qualified as VAI import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM @@ -125,18 +124,16 @@ -- prop_foldl is tested with large array verification problem -prop_maxRight :: Int -> [QC.NonNegative Int] -> QC.Property -prop_maxRight xRef xs_ = - not (null xs_) QC.==> do - -- TODO: remove sort - let xs = VU.modify VAI.sort $ VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ - expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs - res = runST $ do - seg <- Seg.new @_ @(Sum Int) (Seg.recommendedCapacity (VU.length xs) (VU.length xs)) 0 (VU.length xs) - root0 <- Seg.newRoot seg - root <- VU.ifoldM' (Seg.write seg) root0 $ VU.map Sum xs - Seg.maxRight seg root (<= Sum xRef) - in expected QC.=== res +prop_maxRight :: Int -> QC.NonEmptyList (QC.NonNegative Int) -> QC.Property +prop_maxRight xRef (QC.NonEmpty xs_) = + let xs = VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ + expected = VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) xs + res = runST $ do + seg <- Seg.new @_ @(Sum Int) (Seg.recommendedCapacity (VU.length xs) (VU.length xs)) 0 (VU.length xs) + root0 <- Seg.newRoot seg + root <- VU.ifoldM' (Seg.write seg) root0 $ VU.map Sum xs + Seg.maxRight seg root (<= Sum xRef) + in expected QC.=== res tests :: [TestTree] tests = diff --git a/test/Tests/Extra/KdTree.hs b/test/Tests/Extra/KdTree.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/KdTree.hs @@ -0,0 +1,114 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.KdTree where + +import AtCoder.Extra.KdTree qualified as Kt +import Data.Foldable (for_) +import Data.Vector.Algorithms.Intro qualified as VAI +import Data.Vector.Unboxed qualified as VU +import Test.QuickCheck.Monadic as QCM +import Test.Tasty +import Test.Tasty.QuickCheck as QC +import Tests.Util +import Test.Tasty.HUnit +import Debug.Trace + +data Init = Init + { n :: {-# UNPACK #-} !Int, + q :: {-# UNPACK #-} !Int, + refVec :: !(VU.Vector (Int, Int)), + kt :: !Kt.KdTree + } + +rng :: (Int, Int) +rng = (-rngI, rngI) + +rngI :: Int +rngI = 10 + +instance Show Init where + show Init {..} = show n + +instance QC.Arbitrary Init where + arbitrary = do + -- n <- QC.chooseInt (1, 256) + n <- QC.chooseInt (1, 8) + -- q <- QC.chooseInt (1, 256) + q <- QC.chooseInt (1, 4) + refVec <- (VU.fromList <$>) $ QC.vectorOf n $ do + x <- QC.chooseInt rng + y <- QC.chooseInt rng + pure (x, y) + let !_ = traceShow refVec () + let kt = Kt.build2 refVec + pure Init {..} + +data Query + = FindPointsIn ((Int, Int), (Int, Int)) + | FindNearestPoint (Int, Int) + deriving (Show) + +-- | Arbitrary return type for the `Query` result. +data Result + = MI !(Maybe Int) + | V !(VU.Vector Int) + deriving (Show, Eq) + +queryGen :: Int -> QC.Gen Query +queryGen _n = do + QC.oneof + [ FindPointsIn <$> r, + FindNearestPoint <$> xy + ] + where + r = (,) <$> intervalGen' (-rngI) rngI <*> intervalGen' (-rngI) rngI + xy = intervalGen' (-rngI) rngI + +-- | containers. (referencial implementation) +handleRef :: VU.Vector (Int, Int) -> Query -> Result +handleRef vec q = case q of + FindPointsIn ((!x1, !x2), (!y1, !y2)) -> do + V $ VU.findIndices (\(!x, !y) -> x1 <= x && x < x2 && y1 <= y && y < y2) vec + FindNearestPoint (!x, !y) + | VU.null vec -> MI Nothing + | otherwise -> MI + . Just . VU.minIndex + $ VU.map (\(!x', !y') -> (x - x') * (x - x') + (y - y') * (y - y')) vec + +-- | ACL +handleAcl :: Kt.KdTree -> Query -> Result +handleAcl kt q = case q of + FindPointsIn ((!x1, !x2), (!y1, !y2)) -> + V $ Kt.findPointsIn kt x1 x2 y1 y2 (Kt.nKt kt) + FindNearestPoint (!x, !y) -> + MI $ Kt.findNearestPoint kt x y + +prop_randomTest :: Init -> QC.Property +prop_randomTest Init {..} = QCM.monadicIO $ do + qs <- QCM.pick $ QC.vectorOf q (queryGen n) + for_ qs $ \query -> do + let expected = handleRef refVec query + let res = handleAcl kt query + QCM.assertWith (p query expected res) $ show (query, expected, res) + where + p (FindNearestPoint (!x, !y)) (MI (Just a)) (MI (Just b)) = + let (!x1, !y1) = refVec VU.! a + (!x2, !y2) = refVec VU.! b + d1 = (x1 - x) * (x1 - x) + (y1 - y) * (y1 - y) + d2 = (x2 - x) * (x2 - x) + (y2 - y) * (y2 - y) + in a == b || d1 == d2 + p (FindPointsIn _) (V a) (V b) = VU.modify VAI.sort a == VU.modify VAI.sort b + p _ a b = a == b + +unit_zero :: TestTree +unit_zero = testCase "zero" $ do + let kt = Kt.build VU.empty VU.empty + (@?= VU.empty) $ Kt.findPointsIn kt 0 10 0 10 5 + (@?= Nothing) $ Kt.findNearestPoint kt 10 10 + pure () + +tests :: [TestTree] +tests = + [ QC.testProperty "randomTest" prop_randomTest, + unit_zero + ] diff --git a/test/Tests/Extra/LazyKdTree.hs b/test/Tests/Extra/LazyKdTree.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Extra/LazyKdTree.hs @@ -0,0 +1,161 @@ +{-# LANGUAGE RecordWildCards #-} + +module Tests.Extra.LazyKdTree where + +import AtCoder.Extra.LazyKdTree qualified as Lkt +import AtCoder.Extra.Monoid.Affine1 (Affine1 (..)) +import AtCoder.Extra.Monoid.Affine1 qualified as Affine1 +import AtCoder.LazySegTree (segAct) +import AtCoder.ModInt qualified as M +import Control.Monad (when) +import Control.Monad.ST (RealWorld) +import Data.Foldable (for_) +import Data.Semigroup (Sum (..)) +import Data.Vector.Generic.Mutable qualified as VGM +import Data.Vector.Unboxed qualified as VU +import Data.Vector.Unboxed.Mutable qualified as VUM +import Test.QuickCheck.Monadic as QCM +import Test.Tasty +import Test.Tasty.QuickCheck as QC +import Tests.Util +import Test.Tasty.HUnit + +type Mint = M.ModInt998244353 + +modInt :: Int -> Mint +modInt = M.new + +data Init = Init + { n :: {-# UNPACK #-} !Int, + q :: {-# UNPACK #-} !Int, + refM :: !(IO (VUM.MVector RealWorld (Int, Int, Sum Mint))), + ktM :: !(IO (Lkt.LazyKdTree RealWorld (Affine1 Mint) (Sum Mint))) + } + +rng :: (Int, Int) +rng = (-rngI, rngI) + +rngI :: Int +rngI = 10 + +instance Show Init where + show Init {..} = show n + +instance QC.Arbitrary Init where + arbitrary = do + n <- QC.chooseInt (1, 256) + q <- QC.chooseInt (1, 256) + xyws <- (VU.fromList <$>) $ QC.vectorOf n $ do + x <- QC.chooseInt rng + y <- QC.chooseInt rng + w <- Sum . M.new <$> QC.arbitrary @Int + pure (x, y, w) + let refM = VU.thaw xyws + let ktM = Lkt.build3 xyws + pure Init {..} + +data Query + = Write !Int !(Sum Mint) + | Modify !Int !(Sum Mint) + | ModifyM !Int !(Sum Mint) + | Affine !((Int, Int), (Int, Int)) !(Affine1 Mint) + | Prod !((Int, Int), (Int, Int)) + | AllProd + deriving (Show) + +-- | Arbitrary return type for the `Query` result. +data Result + = None + | S !(Sum Mint) + deriving (Show, Eq) + +queryGen :: Int -> QC.Gen Query +queryGen n = do + QC.oneof + [ Write <$> i <*> v, + Modify <$> i <*> v, + ModifyM <$> i <*> v, + Affine <$> r <*> f, + Prod <$> r, + pure AllProd + ] + where + i = QC.chooseInt (0, n - 1) + r = (,) <$> intervalGen' (-rngI) rngI <*> intervalGen' (-rngI) rngI + v = Sum . modInt <$> QC.arbitrary + f = Affine1.new <$> (modInt <$> QC.arbitrary) <*> (modInt <$> QC.arbitrary) + +-- | containers. (referencial implementation) +handleRef :: VUM.MVector RealWorld (Int, Int, Sum Mint) -> Query -> IO Result +handleRef vec q = case q of + Write i v -> do + VGM.write ws i v + pure None + Modify k v -> do + VGM.modify ws (+ v) k + pure None + ModifyM k v -> do + VGM.modify ws (+ v) k + pure None + Affine ((!xl, !xr), (!yl, !yr)) f -> do + xys <- VU.unsafeFreeze $ VUM.zip xs ys + VU.iforM_ xys $ \i (!x, !y) -> do + when (xl <= x && x < xr && yl <= y && y < yr) $ do + VGM.modify ws (segAct f) i + pure None + Prod ((!xl, !xr), (!yl, !yr)) -> do + S + . VU.foldl' + ( \acc (!x, !y, !v) -> + if xl <= x && x < xr && yl <= y && y < yr then acc <> v else acc + ) + mempty + <$> VU.unsafeFreeze vec + AllProd -> do + S <$> VGM.foldl' (<>) mempty ws + where + (!xs, !ys, !ws) = VUM.unzip3 vec + +-- | ac-library-hs. +handleAcl :: (HasCallStack) => Lkt.LazyKdTree RealWorld (Affine1 Mint) (Sum Mint) -> Query -> IO Result +handleAcl kt q = case q of + Write i v -> do + Lkt.write kt i v + pure None + Modify i v -> do + Lkt.modify kt (+ v) i + pure None + ModifyM i v -> do + Lkt.modifyM kt (pure . (+ v)) i + pure None + Affine ((!xl, !xr), (!yl, !yr)) f -> do + Lkt.applyIn kt xl xr yl yr f + pure None + Prod ((!xl, !xr), (!yl, !yr)) -> do + S <$> Lkt.prod kt xl xr yl yr + AllProd -> do + S <$> Lkt.allProd kt + +prop_randomTest :: Init -> QC.Property +prop_randomTest Init {..} = QCM.monadicIO $ do + qs <- QCM.pick $ QC.vectorOf q (queryGen n) + vec <- QCM.run refM + kt <- QCM.run ktM + for_ qs $ \query -> do + expected <- QCM.run $ handleRef vec query + res <- QCM.run $ handleAcl kt query + QCM.assertWith (expected == res) $ show (query, expected, res) + +unit_zero :: TestTree +unit_zero = testCase "zero" $ do + kt <- Lkt.build @_ @(Affine1 Mint) @(Sum Mint) VU.empty VU.empty VU.empty + Lkt.applyIn kt 0 10 0 10 $ Affine1.new 1 1 + (@?= mempty) =<< Lkt.prod kt 0 10 0 10 + (@?= mempty) =<< Lkt.allProd kt + pure () + +tests :: [TestTree] +tests = + [ QC.testProperty "randomTest" prop_randomTest, + unit_zero + ] diff --git a/test/Tests/Extra/Monoid.hs b/test/Tests/Extra/Monoid.hs --- a/test/Tests/Extra/Monoid.hs +++ b/test/Tests/Extra/Monoid.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE DataKinds #-} + module Tests.Extra.Monoid (tests) where import AtCoder.Extra.Monoid diff --git a/test/Tests/Extra/MultiSet.hs b/test/Tests/Extra/MultiSet.hs --- a/test/Tests/Extra/MultiSet.hs +++ b/test/Tests/Extra/MultiSet.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE RecordWildCards #-} module Tests.Extra.MultiSet (tests) where @@ -203,8 +205,7 @@ QCM.assertWith (sizeE == size) $ show ("- size", sizeE, size) pure (ref', recordKey keys query) - else - pure (ref, keys) + else pure (ref, keys) ) (ref0, IS.empty) qs diff --git a/test/Tests/Extra/Semigroup/Matrix.hs b/test/Tests/Extra/Semigroup/Matrix.hs --- a/test/Tests/Extra/Semigroup/Matrix.hs +++ b/test/Tests/Extra/Semigroup/Matrix.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE DataKinds #-} + module Tests.Extra.Semigroup.Matrix (tests) where import AtCoder.Extra.Semigroup.Matrix qualified as Mat diff --git a/test/Tests/FenwickTree.hs b/test/Tests/FenwickTree.hs --- a/test/Tests/FenwickTree.hs +++ b/test/Tests/FenwickTree.hs @@ -2,13 +2,16 @@ module Tests.FenwickTree (tests) where import AtCoder.FenwickTree qualified as FT +import Control.Monad.ST (runST) import Data.Foldable import Data.Vector.Unboxed qualified as VU import System.IO.Unsafe (unsafePerformIO) import Test.Hspec +import Test.QuickCheck qualified as QC import Test.Tasty import Test.Tasty.HUnit import Test.Tasty.Hspec +import Test.Tasty.QuickCheck qualified as QC -- empty -- assign @@ -79,10 +82,32 @@ it "throws error" $ do FT.sum s 5 3 `shouldThrow` anyException +prop_maxRight :: QC.NonNegative Int -> QC.NonEmptyList (QC.NonNegative Int) -> QC.Gen QC.Property +prop_maxRight (QC.NonNegative xRef) (QC.NonEmpty xs_) = do + l0 <- QC.chooseInt (0, length xs_) + let xs = VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ + expected = (l0 +) . VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) $ VU.drop l0 xs + res = runST $ do + ft <- FT.build xs + FT.maxRight ft l0 (<= xRef) + pure $ expected QC.=== res + +prop_minLeft :: QC.NonNegative Int -> QC.NonEmptyList (QC.NonNegative Int) -> QC.Gen QC.Property +prop_minLeft (QC.NonNegative xRef) (QC.NonEmpty xs_) = do + r0 <- QC.chooseInt (0, length xs_) + let xs = VU.fromList $ map (\(QC.NonNegative x) -> x) xs_ + expected = (r0 -) . VU.length . VU.takeWhile (<= xRef) $ VU.scanl1' (+) $ VU.reverse $ VU.take r0 xs + res = runST $ do + ft <- FT.build xs + FT.minLeft ft r0 (<= xRef) + pure $ expected QC.=== res + tests :: [TestTree] tests = [ unit_zero {- overFlowInt -}, unit_naive, unsafePerformIO spec_invalid, - unit_sumMaybeBounds + unit_sumMaybeBounds, + QC.testProperty "maxRight" prop_maxRight, + QC.testProperty "minLeft" prop_minLeft ] diff --git a/test/Tests/ModInt.hs b/test/Tests/ModInt.hs --- a/test/Tests/ModInt.hs +++ b/test/Tests/ModInt.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE DataKinds #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE TypeFamilies #-} diff --git a/test/Tests/Util.hs b/test/Tests/Util.hs --- a/test/Tests/Util.hs +++ b/test/Tests/Util.hs @@ -1,4 +1,4 @@ -module Tests.Util (myForAllShrink, laws, intervalGen) where +module Tests.Util (myForAllShrink, laws, intervalGen, intervalGen') where import Data.Proxy (Proxy (..)) import Data.Typeable (Typeable, typeRep) @@ -43,7 +43,11 @@ -- | Returns an interval [l, r) in [0, n) intervalGen :: Int -> QC.Gen (Int, Int) -intervalGen n = do - l <- QC.chooseInt (0, n) - r <- QC.chooseInt (l, n) +intervalGen = intervalGen' 0 + +-- | Returns an interval [l, r) in [0, n) +intervalGen' :: Int -> Int -> QC.Gen (Int, Int) +intervalGen' xl xr = do + l <- QC.chooseInt (xl, xr) + r <- QC.chooseInt (l, xr) pure (l, r)