ac-library-hs 1.2.2.1 → 1.2.3.0

raw patch · 16 files changed

+1631/−81 lines, 16 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• CHANGELOG.md +4/−0
• ac-library-hs.cabal +6/−1
• src/AtCoder/Extra/KdTree.hs +7/−7
• src/AtCoder/Extra/LazyKdTree.hs +32/−13
• src/AtCoder/Extra/SegTree2d.hs +478/−0
• src/AtCoder/Extra/SegTree2d/Dense.hs +326/−0
• src/AtCoder/Extra/Vector.hs +67/−0
• src/AtCoder/Extra/WaveletMatrix.hs +9/−9
• src/AtCoder/Extra/WaveletMatrix/Raw.hs +12/−12
• src/AtCoder/Extra/WaveletMatrix2d.hs +7/−3
• src/AtCoder/FenwickTree.hs +0/−7
• src/AtCoder/Internal/Assert.hs +99/−6
• src/AtCoder/Internal/Queue.hs +287/−23
• test/Main.hs +4/−0
• test/Tests/Extra/SegTree2d.hs +147/−0
• test/Tests/Extra/SegTree2d/Dense.hs +146/−0

@@ -1,5 +1,9 @@ # Revision history for acl-hs +## 1.2.3.0 -- March 2025++- Added `Extra.SegTree2d` and `Extra.SegTree2d.Dense`.+ ## 1.2.2.1 -- March 2025  - Reduced build time with `ST` monad and `INLINEABLE` pragmas.

@@ -4,7 +4,7 @@ -- PVP summary:  +-+------- breaking API changes --               | | +----- non-breaking API additions --               | | | +--- code changes with no API change-version:         1.2.2.1+version:         1.2.3.0 synopsis:        Data structures and algorithms description:   Haskell port of [ac-library](https://github.com/atcoder/ac-library), a library for competitive@@ -79,6 +79,8 @@     AtCoder.Extra.MultiSet     AtCoder.Extra.Pdsu     AtCoder.Extra.Pool+    AtCoder.Extra.SegTree2d+    AtCoder.Extra.SegTree2d.Dense     AtCoder.Extra.Semigroup.Matrix     AtCoder.Extra.Semigroup.Permutation     AtCoder.Extra.Seq@@ -88,6 +90,7 @@     AtCoder.Extra.Tree.Hld     AtCoder.Extra.Tree.Lct     AtCoder.Extra.Tree.TreeMonoid+    AtCoder.Extra.Vector     AtCoder.Extra.WaveletMatrix     AtCoder.Extra.WaveletMatrix.BitVector     AtCoder.Extra.WaveletMatrix.Raw@@ -150,6 +153,8 @@     Tests.Extra.Math     Tests.Extra.Monoid     Tests.Extra.MultiSet+    Tests.Extra.SegTree2d+    Tests.Extra.SegTree2d.Dense     Tests.Extra.Semigroup.Matrix     Tests.Extra.Semigroup.Permutation     Tests.Extra.Seq

@@ -152,7 +152,7 @@   where     (!xs, !ys) = VU.unzip xys --- | \(O(n \log n)\) Collects points in \([x_l, x_r) \times [y_l, y_r)\).+-- | \(O(n)\) Collects points in \([x_1, x_2) \times [y_1, y_2)\). -- -- @since 1.2.2.0 {-# INLINEABLE findPointsIn #-}@@ -160,17 +160,17 @@   (HasCallStack) =>   -- | `KdTree`   KdTree ->-  -- | \(x_l\)+  -- | \(x_1\)   Int ->-  -- | \(x_r\)+  -- | \(x_2\)   Int ->-  -- | \(y_l\)+  -- | \(y_1\)   Int ->-  -- | \(y_r\)+  -- | \(y_2\)   Int ->-  -- | Maximum number of points in \([x_l, x_r) \times [y_l, y_r)\).+  -- | Maximum number of points in \([x_1, x_2) \times [y_1, y_2)\).   Int ->-  -- | Point indices in \([x_l, x_r) \times [y_l, y_r)\).+  -- | Point indices in \([x_1, x_2) \times [y_1, y_2)\).   VU.Vector Int findPointsIn KdTree {..} x1 x2 y1 y2 capacity   | nKt == 0 = VU.empty

@@ -36,6 +36,7 @@     SegAct (..),      -- * Constructors+    new,     build,     build2,     build3,@@ -106,6 +107,23 @@     posLkt :: !(VU.Vector Int)   } +-- | \(O(n \log n)\) Creates a `LazyKdTree` from @xs@ and @ys@.+--+-- ==== Constraints+-- - \(|\mathrm{xs}| = |\mathrm{ys}|+--+-- @since 1.2.3.0+{-# INLINE new #-}+new ::+  (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>+  -- | \(x\) coordnates+  VU.Vector Int ->+  -- | \(y\) coordnates+  VU.Vector Int ->+  -- | `LazyKdTree`+  m (LazyKdTree (PrimState m) f a)+new xs ys = stToPrim $ buildST xs ys (VU.replicate (VU.length xs) mempty)+ -- | \(O(n \log n)\) Creates a `LazyKdTree` from @xs@, @ys@ and @ws@ vectors. -- -- ==== Constraints@@ -205,6 +223,7 @@   -- | Monadic tuple   m () modifyM kt@LazyKdTree {..} f i0 = do+  let !_ = ACIA.checkIndex "AtCoder.Extra.LazyKdTree.modifyM" i0 nLkt   let i_ = posLkt VG.! i0   -- propagate lazily propagated monoid actions from the root:   stToPrim $ for_ [logLkt, logLkt - 1 .. 1] $ \k -> do@@ -221,7 +240,7 @@             inner i'   stToPrim $ inner i_ --- | \(O(\log n)\) Returns monoid product in \([x_l, x_r) \times [y_l, y_r)\).+-- | \(O(\sqrt n)\) Returns monoid product \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\). -- -- @since 1.2.2.0 {-# INLINE prod #-}@@ -229,19 +248,19 @@   (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\)+  -- | \(x_1\)   Int ->-  -- | \(x_r\)+  -- | \(x_2\)   Int ->-  -- | \(y_l\)+  -- | \(y_1\)   Int ->-  -- | \(y_r\)+  -- | \(y_2\)   Int ->-  -- | Monoid product in \([x_l, x_r) \times [y_l, y_r)\)+  -- | \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\)   m a prod kt x1 x2 y1 y2 = stToPrim $ prodST kt x1 x2 y1 y2 --- | \(O(1)\) Returns monoid product of all the points.+-- | \(O(1)\) Returns monoid product \(\Pi_{p \in [-\infty, \infty) \times [-\infty, \infty)} a_p\). -- -- @since 1.2.2.0 {-# INLINE allProd #-}@@ -249,13 +268,13 @@   (PrimMonad m, Monoid a, VU.Unbox a) =>   -- | `LazyKdTree`   LazyKdTree (PrimState m) f a ->-  -- | Monoid product in the whole space.+  -- | \(\Pi_{p \in [-\infty, \infty) \times [-\infty, \infty)} a_p\)   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)\).+-- | \(O(\log n)\) Applies a monoid action to points in \([x_1, x_2) \times [y_1, y_2)\). -- -- @since 1.2.2.0 {-# INLINE applyIn #-}@@ -263,13 +282,13 @@   (HasCallStack, PrimMonad m, Eq f, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) =>   -- | `LazyKdTree`   LazyKdTree (PrimState m) f a ->-  -- | \(x_l\)+  -- | \(x_1\)   Int ->-  -- | \(x_r\)+  -- | \(x_2\)   Int ->-  -- | \(y_l\)+  -- | \(y_1\)   Int ->-  -- | \(y_r\)+  -- | \(y_2\)   Int ->   -- | \(f\)   f ->

@@ -0,0 +1,478 @@+{-# LANGUAGE RecordWildCards #-}++-- | Two-dimensional segment tree for commutative monoids at fixed points.+--+-- ==== SegTree2d vs WaveletMatrix2d+-- They basically the same functionalities and performance, however, in @ac-library-hs@, `SegTree2d`+-- has better API and even outperforms @WaveletMatrix2d@.+--+-- ==== __Examples__+-- Create a two-dimensional segment tree for points \((0, 0)\) with weight \(10\) and \((1, 1)\)+-- with weight \(20\):+--+-- >>> import AtCoder.Extra.SegTree2d qualified as Seg+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> seg <- Seg.build3 @_ @(Sum Int) $ VU.fromList [(0, 0, 10), (1, 1, 20)]+--+-- Get monoid product in \([x_1, x_2) \times [y_1, y_2)\) with `prod`:+--+-- >>> Seg.prod seg {- x -} 0 2 {- y -} 0 2+-- Sum {getSum = 30}+--+-- Monoid values can be altered:+--+-- >>> Seg.write seg 1 50+-- >>> Seg.prod seg {- x -} 1 2 {- y -} 1 2+-- Sum {getSum = 50}+--+-- >>> Seg.allProd seg+-- Sum {getSum = 60}+--+-- >>> Seg.count seg {- x -} 0 2 {- y -} 0 2+-- 2+--+-- @since 1.2.3.0+module AtCoder.Extra.SegTree2d+  ( -- * SegTree2d+    SegTree2d (..),++    -- * Constructors+    new,+    build,+    build2,+    build3,++    -- * Write++    -- read,+    write,+    modify,+    modifyM,++    -- * Monoid product+    prod,+    allProd,++    -- * Count+    count,+  )+where++import AtCoder.Extra.Bisect (lowerBound)+import AtCoder.Extra.Vector (argsort)+import AtCoder.Internal.Assert qualified as ACIA+import AtCoder.Internal.Bit qualified as ACIB+import Control.Monad (when)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST)+import Data.Bits+import Data.Foldable (for_)+import Data.Maybe (fromJust, fromMaybe)+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)+import Prelude hiding (read)++-- | Two-dimensional segment tree.+--+-- @since 1.2.3.0+data SegTree2d s a = SegTree2d+  { -- | The number of nodes.+    --+    -- @since 1.2.3.0+    nSt :: {-# UNPACK #-} !Int,+    -- | The number of distinct \(x\) coordinates.+    --+    -- @since 1.2.3.0+    nxSt :: {-# UNPACK #-} !Int,+    -- | \(\lceil \log_2 (\mathrm{nx} + 1) \rceil\)+    --+    -- @since 1.2.3.0+    logSt :: {-# UNPACK #-} !Int,+    -- | \(2^{\mathrm{logSt}}\)+    --+    -- @since 1.2.3.0+    sizeSt :: {-# UNPACK #-} !Int,+    -- | \(x\) coordinates sorted and uniquified+    --+    -- @since 1.2.3.0+    dictXSt :: {-# UNPACK #-} !(VU.Vector Int),+    -- | \(y\) coordinates sorted (not uniquified)+    --+    -- @since 1.2.3.0+    allYSt :: !(VU.Vector Int),+    -- |+    --+    -- @since 1.2.3.0+    posSt :: !(VU.Vector Int),+    -- |+    --+    -- @since 1.2.3.0+    indptrSt :: !(VU.Vector Int),+    -- | Monoid values+    --+    -- @since 1.2.3.0+    dataSt :: !(VUM.MVector s a),+    -- |+    --+    -- @since 1.2.3.0+    toLeftSt :: !(VU.Vector Int)+  }++-- | \(O(n \log n)\) Creates a `SegTree2d` from a vector of \((x, y)\) coordinates.+--+-- @since 1.2.3.0+{-# INLINEABLE new #-}+new ::+  (PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | \((x, y)\) vector+  VU.Vector (Int, Int) ->+  -- | Two-dimensional segment tree+  m (SegTree2d (PrimState m) a)+new xys = stToPrim $ buildST xs ys (VU.replicate n mempty)+  where+    n = VU.length xys+    (!xs, !ys) = VU.unzip xys++-- | \(O(n \log n)\) Creates a `SegTree2d` from vectors of \(x\), \(y\) and \(w\).+--+-- @since 1.2.3.0+{-# INLINE build #-}+build ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | \(x\) vector+  VU.Vector Int ->+  -- | \(y\) vector+  VU.Vector Int ->+  -- | \(w\) vector+  VU.Vector a ->+  -- | Two-dimensional segment tree+  m (SegTree2d (PrimState m) a)+build xs ys ws = stToPrim $ buildST xs ys ws++-- | \(O(n \log n)\) Creates a `SegTree2d` from vectors of \((x, y)\) and \(w\).+--+-- @since 1.2.3.0+{-# INLINE build2 #-}+build2 ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | \((x, y)\) vector+  VU.Vector (Int, Int) ->+  -- | \(w\) vector+  VU.Vector a ->+  -- | Two-dimensional segment tree+  m (SegTree2d (PrimState m) a)+build2 xys ws = stToPrim $ do+  let (!xs, !ys) = VU.unzip xys+  buildST xs ys ws++-- | \(O(n \log n)\) Creates a `SegTree2d` from a vector of \((x, y, w)\).+--+-- @since 1.2.3.0+{-# INLINE build3 #-}+build3 :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => VU.Vector (Int, Int, a) -> m (SegTree2d (PrimState m) a)+build3 xyws = stToPrim $ do+  let (!xs, !ys, !ws) = VU.unzip3 xyws+  buildST xs ys ws++-- -- | \(O(\log n)\) Read the \(k\)-th original point's monoid value.+-- --+-- -- @since 1.2.3.0+-- {-# INLINE read #-}+-- read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree2d (PrimState m) a -> Int -> m ()+-- read SegTree2d {..} i = do+--   let !_ = ACIA.checkIndex "AtCoder.Extra.SegTree2d.read" i nSt+--   -- FIXME: pos is incorrect+--   VGM.read dataSt (posSt VG.! i)++-- | \(O(\log n)\) Writes to the \(k\)-th original point's monoid value.+--+-- @since 1.2.3.0+{-# INLINE write #-}+write ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Two-dimensional segment tree+  SegTree2d (PrimState m) a ->+  -- | Original point index+  Int ->+  -- | New monoid value+  a ->+  m ()+write seg i x = stToPrim $ do+  modifyM seg (pure . const x) i++-- | \(O(\log n)\) Given \(f\), modofies the \(k\)-th original point's monoid value.+--+-- @since 1.2.3.0+{-# INLINE modify #-}+modify ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Two-dimensional segment tree+  SegTree2d (PrimState m) a ->+  -- | Function that alters the monoid value+  (a -> a) ->+  -- | Original point index+  Int ->+  m ()+modify seg f i = stToPrim $ do+  modifyM seg (pure . f) i++-- | \(O(\log n)\) Given \(f\), modofies the \(k\)-th original point's monoid value.+--+-- @since 1.2.3.0+{-# INLINEABLE modifyM #-}+modifyM ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Two-dimensional segment tree+  SegTree2d (PrimState m) a ->+  -- | Function that alters the monoid value+  (a -> m a) ->+  -- | Original point index+  Int ->+  m ()+modifyM seg@SegTree2d {..} f rawIdx = do+  let !_ = ACIA.checkIndex "AtCoder.Extra.SegTree2d.modifyM" rawIdx nSt+  inner 1 $ posSt VG.! rawIdx+  where+    inner i p = do+      let indptrI = indptrSt VG.! i+      modifyIST seg f i $ p - indptrI+      when (i < sizeSt) $ do+        let lc = toLeftSt VG.! p - toLeftSt VG.! indptrI+        let rc = (p - indptrI) - lc+        if toLeftSt VG.! (p + 1) - toLeftSt VG.! p /= 0+          then do+            let i' = 2 * i + 0+            let p' = indptrSt VG.! i' + lc+            inner i' p'+          else do+            let i' = 2 * i + 1+            let p' = indptrSt VG.! i' + rc+            inner i' p'++-- | \(O(\log^2 n)\) Returns monoid product \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\).+--+-- @since 1.2.3.0+{-# INLINE prod #-}+prod ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Two-dimensional segment tree+  SegTree2d (PrimState m) a ->+  -- | \(x_1\)+  Int ->+  -- | \(x_2\)+  Int ->+  -- | \(y_1\)+  Int ->+  -- | \(y_2\)+  Int ->+  -- | \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\)+  m a+prod seg lx rx ly ry = stToPrim $ prodST seg lx rx ly ry++-- | \(O(1)\) Returns monoid product \(\Pi_{p \in [-\infty, \infty) \times [-\infty, \infty)} a_p\).+--+-- @since 1.2.3.0+{-# INLINE allProd #-}+allProd :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree2d (PrimState m) a -> m a+allProd seg = fromMaybe mempty <$> VGM.readMaybe (dataSt seg) 1++-- | \(O(\log n)\) Returns the number of points in \([x_1, x_2) \times [y_1, y_2)\).+--+-- @since 1.2.3.0+{-# INLINE count #-}+count ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Two-dimensional segment tree+  SegTree2d (PrimState m) a ->+  -- | \(x_1\)+  Int ->+  -- | \(x_2\)+  Int ->+  -- | \(y_1\)+  Int ->+  -- | \(y_2\)+  Int ->+  -- | The number of points in \([x_1, x_2) \times [y_1, y_2)\)+  m Int+count seg lx rx ly ry = stToPrim $ countST seg lx rx ly ry++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE buildST #-}+buildST :: forall s a. (HasCallStack, Monoid a, VU.Unbox a) => VU.Vector Int -> VU.Vector Int -> VU.Vector a -> ST s (SegTree2d s a)+buildST xs ys ws = do+  let nSt = VU.length xs+  let !_ = ACIA.runtimeAssert (nSt == VU.length ys && nSt == VU.length ws) "AtCoder.Extra.SegTree2d.buildST: length mismatch among `xs`, `ys` and `ws`"++  -- TODO: use radix sort?+  -- we don't have to take `uniq` though:+  let dictXSt = VU.uniq $ VU.modify VAI.sort xs+  let nxSt = VU.length dictXSt+  let logSt = countTrailingZeros $ ACIB.bitCeil (nxSt + 1)+  let sizeSt = bit logSt+  let compressedXs = VU.map (fromJust . lowerBound dictXSt) xs++  -- TODO: what is this?+  let indptrSt = VU.create $ do+        indptr <- VUM.replicate (2 * sizeSt + 1) (0 :: Int)+        VU.forM_ compressedXs $ \i -> do+          let inner j+                | j /= 0 = do+                    VGM.modify indptr (+ 1) (j + 1) -- +1 for perfix sum+                    inner (j `div` 2)+                | otherwise = pure ()+          inner $ i + sizeSt+        -- calculate prefix sum in-place:+        VUM.iforM_ (VUM.init indptr) $ \i x -> do+          VGM.modify indptr (+ x) (i + 1)+        pure indptr++  dataSt <- VUM.replicate (2 * VU.last indptrSt) (mempty :: a)++  let yis = argsort ys+  let posSt = VU.create $ do+        vec <- VUM.replicate nSt (0 :: Int)+        VU.iforM_ yis $ \i yi -> do+          -- FIXME: It overwrites duplicate yi to i. Isn't it just `accumulate max`?+          VGM.write vec yi i+        pure vec++  -- +1 for prefix cumulative sum+  toLeftSt <- do+    toLeft <- VUM.replicate (indptrSt VG.! sizeSt + 1) (0 :: Int)+    ptr <- VU.thaw indptrSt+    VU.forM_ yis $ \rawIdx -> do+      let inner i j+            | i > 0 = do+                p <- VGM.read ptr i+                VGM.write ptr i $ p + 1+                VGM.write dataSt (indptrSt VG.! (i + 1) + p) $ ws VG.! rawIdx+                when (j /= -1) $ do+                  VGM.write toLeft (p + 1) $ if even j then 1 else 0+                inner (i `div` 2) i+            | otherwise = pure ()+      let i0 = compressedXs VG.! rawIdx + sizeSt+      inner i0 (-1)++    -- calculate prefix cumulative sum in-place:+    VUM.iforM_ (VUM.init toLeft) $ \i x -> do+      VGM.modify toLeft (+ x) (i + 1)+    VU.unsafeFreeze toLeft++  for_ [0 .. 2 * sizeSt - 1] $ \i -> do+    let off = 2 * indptrSt VG.! i+    let n = indptrSt VG.! (i + 1) - indptrSt VG.! i+    for_ [n - 1, n - 2 .. 1] $ \j -> do+      xl <- VGM.read dataSt $ off + 2 * j + 0+      xr <- VGM.read dataSt $ off + 2 * j + 1+      VGM.write dataSt (off + j) $! xl <> xr++  -- TODO: why not uniquified?+  let allYSt = VU.modify VAI.sort ys+  pure SegTree2d {..}++{-# INLINEABLE modifyIST #-}+modifyIST :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()+modifyIST SegTree2d {..} f i j0 = do+  -- TODO: LID?+  let lid = indptrSt VG.! i+  let off = 2 * lid+  let inner j_ = do+        when (j_ > 1) $ do+          let j = j_ `div` 2+          xl <- VGM.read dataSt (off + 2 * j + 0)+          xr <- VGM.read dataSt (off + 2 * j + 1)+          VGM.write dataSt (off + j) $! xl <> xr+          inner j+  let j = j0 + indptrSt VG.! (i + 1) - lid+  VGM.modifyM dataSt f (off + j)+  stToPrim $ inner j++{-# INLINEABLE prodST #-}+prodST :: forall s a. (HasCallStack, Monoid a, VU.Unbox a) => SegTree2d s a -> Int -> Int -> Int -> Int -> ST s a+prodST seg@SegTree2d {..} lx rx ly ry = do+  let a0 = fromMaybe (VU.length allYSt) $ lowerBound allYSt ly+  let b0 = fromMaybe (VU.length allYSt) $ lowerBound allYSt ry+  dfs mempty 1 0 sizeSt a0 b0+  where+    !_ = ACIA.runtimeAssert (lx <= rx && ly <= ry) "AtCoder.Extra.SegTree2d.prodST: given invalid rectangle"+    !l0 = fromMaybe (VU.length dictXSt) $ lowerBound dictXSt lx+    !r0 = fromMaybe (VU.length dictXSt) $ lowerBound dictXSt rx+    dfs :: a -> Int -> Int -> Int -> Int -> Int -> ST s a+    dfs !res i l r a b+      -- empty rect+      | a == b = pure res+      -- not intersecting+      | r0 <= l || r <= l0 = pure res+      -- fully contained in [l_0, r_0)+      | l0 <= l && r <= r0 = do+          xi <- prodIST seg i a b+          pure $! xi <> res+      | otherwise = do+          let indptrI = indptrSt VG.! i+          let toLeftI = toLeftSt VG.! indptrI+          let la = toLeftSt VG.! (indptrI + a) - toLeftI+          let ra = a - la+          let lb = toLeftSt VG.! (indptrI + b) - toLeftI+          let rb = b - lb+          let m = (l + r) `div` 2+          res' <- dfs res (2 * i + 0) l m la lb+          dfs res' (2 * i + 1) m r ra rb++{-# INLINEABLE prodIST #-}+prodIST :: (HasCallStack, Monoid a, VU.Unbox a) => SegTree2d s a -> Int -> Int -> Int -> ST s a+prodIST SegTree2d {..} i a b = inner mempty (n + a) (n + b - 1)+  where+    lid = indptrSt VG.! i+    off = 2 * lid+    n = indptrSt VG.! (i + 1) - lid+    -- NOTE: we're using inclusive interval [l, r] for simplicity+    inner !res l r+      | l <= r = do+          res' <-+            if testBit l 0+              then (res <>) <$> VGM.read dataSt (off + l)+              else pure res+          res'' <-+            if not $ testBit r 0+              then (<> res') <$> VGM.read dataSt (off + r)+              else pure res'+          inner res'' ((l + 1) .>>. 1) ((r - 1) .>>. 1)+      | otherwise = pure res++{-# INLINEABLE countST #-}+countST :: forall s a. (HasCallStack, Monoid a, VU.Unbox a) => SegTree2d s a -> Int -> Int -> Int -> Int -> ST s Int+countST SegTree2d {..} lx rx ly ry = do+  let a0 = fromMaybe (VU.length allYSt) $ lowerBound allYSt ly+  let b0 = fromMaybe (VU.length allYSt) $ lowerBound allYSt ry+  dfs 0 1 0 sizeSt a0 b0+  where+    !_ = ACIA.runtimeAssert (lx <= rx && ly <= ry) "AtCoder.Extra.SegTree2d.countST: given invalid rectangle"+    !l0 = fromMaybe (VU.length dictXSt) $ lowerBound dictXSt lx+    !r0 = fromMaybe (VU.length dictXSt) $ lowerBound dictXSt rx+    dfs :: Int -> Int -> Int -> Int -> Int -> Int -> ST s Int+    dfs (res :: Int) i l r a b+      -- empty rect+      | a == b = pure res+      -- not intersecting+      | r0 <= l || r <= l0 = pure res+      -- fully contained in [l_0, r_0)+      | l0 <= l && r <= r0 = do+          pure $! res + b - a+      | otherwise = do+          let indptrI = indptrSt VG.! i+          let toLeftI = toLeftSt VG.! indptrI+          let la = toLeftSt VG.! (indptrI + a) - toLeftI+          let ra = a - la+          let lb = toLeftSt VG.! (indptrI + b) - toLeftI+          let rb = b - lb+          let m = (l + r) `div` 2+          res' <- dfs res (2 * i + 0) l m la lb+          dfs res' (2 * i + 1) m r ra rb

@@ -0,0 +1,326 @@+{-# LANGUAGE RecordWildCards #-}++-- | Two-dimensional segment tree for commutative monoids in \([0, w) \times [0, h)\).+--+-- ==== __Internals__+-- Take a 2x4 matrix as an example:+--+-- @+-- 5 6 7 8+-- 1 2 3 4+-- @+--+-- Extend each row as a segment tree:+--+-- @+--  - 22 11 15  5  6  7  8+--  - 10  3  7  1  2  3  4+-- @+--+-- Then extend each column as a segment tree:+--+-- @+--  -  -  -  -  -  -  -  -+--  - 30 14 22  6  8 10 12+--  - 26 11 15  5  6  7  8+--  - 10  3  7  1  2  3  4+-- @+--+-- ==== __ Example__+-- Create a two-dimensional segment tree for size (w, h) = (4, 2):+--+-- >>> import AtCoder.Extra.SegTree2d.Dense qualified as Seg+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> seg <- Seg.build @_ @(Sum Int) 4 2 $ VU.fromList [0, 1, 2, 3, 4, 5, 6, 7]+--+-- Get monoid product in \([x_1, x_2) \times [y_1, y_2)\) with `prod`:+--+-- >>> Seg.prod seg {- x -} 1 4 {- y -} 0 1+-- Sum {getSum = 6}+--+-- Monoid values can be altered:+--+-- >>> Seg.write seg 1 1 20+-- >>> Seg.prod seg {- x -} 0 2 {- y -} 0 2+-- Sum {getSum = 25}+--+-- >>> Seg.allProd seg+-- Sum {getSum = 43}+--+-- @since 1.2.3.0+module AtCoder.Extra.SegTree2d.Dense+  ( -- * DenseSegTree2d+    DenseSegTree2d (..),++    -- * Constructors+    new,+    build,+    build',++    -- * Read+    read,+    readMaybe,++    -- * Write+    write,+    modify,+    modifyM,++    -- * Monoid product+    prod,+    allProd,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import AtCoder.Internal.Bit qualified as ACIB+import Control.Monad (when)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST)+import Data.Bits+import Data.Foldable (for_)+import Data.Maybe (fromJust, fromMaybe)+import Data.Vector qualified as V+import Data.Vector.Algorithms.Intro qualified as VAI+import Data.Vector.Generic qualified as VG+import Data.Vector.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)++-- | Two-dimensional segment tree.+--+-- @since 1.2.3.0+data DenseSegTree2d s a = DenseSegTree2d+  { -- | Height+    --+    -- @since 1.2.3.0+    hDst :: {-# UNPACK #-} !Int,+    -- | Width+    --+    -- @since 1.2.3.0+    wDst :: {-# UNPACK #-} !Int,+    -- | Monoid values+    --+    -- @since 1.2.3.0+    dataDst :: !(VUM.MVector s a)+  }++-- | \(O(hw)\) Creates a `DenseSegTree2d` for \([0, w) \times [0, h)\) from \(w\) and \(h\).+--+-- @since 1.2.3.0+{-# INLINEABLE new #-}+new ::+  (PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Width+  Int ->+  -- | Height+  Int ->+  -- | Dense, two-dimensional segment tree+  m (DenseSegTree2d (PrimState m) a)+new wDst hDst = stToPrim $ do+  dataDst <- VUM.replicate (4 * wDst * hDst) mempty+  pure DenseSegTree2d {..}++-- | \(O(hw)\) Creates a `DenseSegTree2d` from width, height and one-dimensional vector of+-- monoid values.+--+-- @since 1.2.3.0+{-# INLINE build #-}+build ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Width+  Int ->+  -- | Height+  Int ->+  -- | Vector of monoid values+  VU.Vector a ->+  -- | Dense, two-dimensional segment tree+  m (DenseSegTree2d (PrimState m) a)+build w h xs = stToPrim $ buildST $ V.unfoldrExactN h (VU.splitAt w) xs+  where+    !_ = ACIA.runtimeAssert (VU.length xs == w * h) "AtCoder.Extra.SegTree2d.Dense.build: vector length mismatch"++-- | \(O(hw)\) Creates a `DenseSegTree2d` from a two-dimensional vector of monoid values.+-- The vector must be indexed by \(y\) first then \(x\): @vec V.! y VU.! x@.+--+-- ==== Constraints+-- - The length of the monoid value vector must be \(hw\).+--+-- @since 1.2.3.0+{-# INLINE build' #-}+build' ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Two-dimensional vector of monoid values+  V.Vector (VU.Vector a) ->+  -- | Dense, two-dimensional segment tree+  m (DenseSegTree2d (PrimState m) a)+build' xs = stToPrim $ buildST xs++-- | \(O(1)\) Returns the monoid value at \((x, y)\).+--+-- @since 1.2.3.0+{-# INLINE read #-}+read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> m a+read seg@DenseSegTree2d {..} x y = do+  let !_ = ACIA.checkPoint2d "AtCoder.Extra.SegTree2d.Dense.read" x y wDst hDst+  VGM.read dataDst $ idx wDst (y + hDst) (x + wDst)++-- | \(O(1)\) Returns the monoid value at \((x, y)\), or `Nothing` if the point is out of the+-- bounds.+--+-- @since 1.2.3.0+{-# INLINE readMaybe #-}+readMaybe :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> m (Maybe a)+readMaybe seg@DenseSegTree2d {..} x y+  | ACIA.testPoint2d x y wDst hDst = do+      Just <$> VGM.read dataDst (idx wDst (y + hDst) (x + wDst))+  | otherwise = pure Nothing++-- | \(O(\log  h \log w)\) Writes to the \(k\)-th original point's monoid value.+--+-- @since 1.2.3.0+{-# INLINE write #-}+write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> a -> m ()+write seg@DenseSegTree2d {..} x y a = stToPrim $ do+  modifyM seg (pure . const a) x y++-- | \(O(\log  h \log w)\) Given \(f\), modofies the monoid value at \((x, y)\).+--+-- @since 1.2.3.0+{-# INLINE modify #-}+modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> (a -> a) -> Int -> Int -> m ()+modify seg f x y = stToPrim $ do+  modifyM seg (pure . f) x y++-- | \(O(\log h \log w)\) Given \(f\), modofies the monoid value at \((x, y)\).+--+-- @since 1.2.3.0+{-# INLINEABLE modifyM #-}+modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()+modifyM DenseSegTree2d {..} f x0_ y0_ = do+  let !_ = ACIA.checkPoint2d "AtCoder.Extra.SegTree2d.Dense.modifyM" x0_ y0_ wDst hDst+  let y0 = y0_ + hDst+  let x0 = x0_ + wDst+  VGM.modifyM dataDst f (idx wDst y0 x0)+  stToPrim $ do+    -- right to left+    let updateCurrentRow 0 = pure ()+        updateCurrentRow x = do+          xl <- VGM.read dataDst (idx wDst y0 (2 * x + 0))+          xr <- VGM.read dataDst (idx wDst y0 (2 * x + 1))+          VGM.write dataDst (idx wDst y0 x) $! xl <> xr+          updateCurrentRow (x `div` 2)+    updateCurrentRow (x0 `div` 2)++    -- down to up+    let updateOtherRow 0 = pure ()+        updateOtherRow y = do+          let updateRow 0 = pure ()+              updateRow x = do+                xl <- VGM.read dataDst (idx wDst (2 * y + 0) x)+                xr <- VGM.read dataDst (idx wDst (2 * y + 1) x)+                VGM.write dataDst (idx wDst y x) $! xl <> xr+                updateRow (x `div` 2)+          updateRow x0+          updateOtherRow (y `div` 2)+    updateOtherRow (y0 `div` 2)++-- | \(O(\log h \log w)\) Returns monoid product \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\).+--+-- @since 1.2.3.0+{-# INLINE prod #-}+prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a+prod seg@DenseSegTree2d {..} x1 x2 y1 y2 = stToPrim $ do+  let !_ = ACIA.checkRectShape "AtCoder.Extra.SegTree2d.Dense.prodST" x1 x2 y1 y2+  prodST seg (max 0 x1) (min wDst x2) (max 0 y1) (min hDst y2)++-- | \(O(1)\) Returns monoid product \(\Pi_{p \in [0, w) \times [0, h)} a_p\).+--+-- @since 1.2.3.0+{-# INLINE allProd #-}+allProd :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> m a+allProd DenseSegTree2d {..} = stToPrim $ do+  -- FIXME: correct?+  fromMaybe mempty <$> VGM.readMaybe dataDst (idx wDst 1 1)++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINE idx #-}+idx :: Int -> Int -> Int -> Int+idx w y x = y * (2 * w) + x++{-# INLINEABLE buildST #-}+buildST :: (HasCallStack, Monoid a, VU.Unbox a) => V.Vector (VU.Vector a) -> ST s (DenseSegTree2d s a)+buildST vec = do+  let hDst = V.length vec+  let wDst = VU.length (V.head vec)++  -- NOTE: It's zero-based and we do not ceil H/W to 2^n, still the indexing works fine:+  --           1  2  3+  -- 11  6  5  1  2  3+  dataDst <- VUM.replicate (4 * hDst * wDst) mempty++  -- copy the base matrix in [w, 2w) \times [h, 2h):+  V.iforM_ vec $ \y vs -> do+    VU.iforM_ vs $ \x v -> do+      VGM.write dataDst (idx wDst (hDst + y) (wDst + x)) v++  -- extend the row (y >= h) as a segment tree's internal vector:+  for_ [hDst .. 2 * hDst - 1] $ \y -> do+    for_ [wDst - 1, wDst - 2 .. 0] $ \x -> do+      xl <- VGM.read dataDst (idx wDst y (2 * x + 0))+      xr <- VGM.read dataDst (idx wDst y (2 * x + 1))+      VGM.write dataDst (idx wDst y x) $! xl <> xr++  -- extend each column as a segment tree:+  -- NOTE (pref): interate from y then x for contiguous memory access+  for_ [hDst - 1, hDst - 2 .. 0] $ \y -> do+    for_ [0 .. 2 * wDst - 1] $ \x -> do+      xl <- VGM.read dataDst (idx wDst (2 * y + 0) x)+      xr <- VGM.read dataDst (idx wDst (2 * y + 1) x)+      VGM.write dataDst (idx wDst y x) $! xl <> xr++  pure DenseSegTree2d {..}++{-# INLINEABLE prodST #-}+prodST :: (HasCallStack, Monoid a, VU.Unbox a) => DenseSegTree2d s a -> Int -> Int -> Int -> Int -> ST s a+prodST seg@DenseSegTree2d {..} x1 x2 y1 y2 = do+  inner mempty (y1 + hDst) (y2 + hDst - 1)+  where+    -- inclusive interval [yl, yr]+    inner !acc yl yr+      | yl > yr = pure acc+      | otherwise = do+          acc' <-+            if testBit yl 0+              then (acc <>) <$> prodX seg yl x1 x2+              else pure acc+          acc'' <-+            if not $ testBit yr 0+              then (<> acc') <$> prodX seg yr x1 x2+              else pure acc'+          inner acc'' ((yl + 1) .>>. 1) ((yr - 1) .>>. 1)++{-# INLINEABLE prodX #-}+prodX :: (HasCallStack, Monoid a, VU.Unbox a) => DenseSegTree2d s a -> Int -> Int -> Int -> ST s a+prodX DenseSegTree2d {..} y x1 x2 = do+  inner mempty (x1 + wDst) (x2 + wDst - 1)+  where+    -- inclusive interval [xl, xr]+    inner !acc xl xr+      | xl > xr = pure acc+      | otherwise =do+          acc' <-+            if testBit xl 0+              then (acc <>) <$> VGM.read dataDst (idx wDst y xl)+              else pure acc+          acc'' <-+            if not $ testBit xr 0+              then (<> acc') <$> VGM.read dataDst (idx wDst y xr)+              else pure acc'+          inner acc'' ((xl + 1) .>>. 1) ((xr - 1) .>>. 1)

@@ -0,0 +1,67 @@+-- | Miscellaneous vector methods.+--+-- @since 1.2.2.0+module AtCoder.Extra.Vector+  ( argsort,+    unsafePermuteInPlace,+    unsafePermuteInPlaceST,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import Control.Monad (unless)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST)+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++-- TODO: test `unsafePermuteInPlace`+-- TODO: is `unsafePermuteInPlace` fast enough as specialized one?++-- | \(O(n \log n)\) Returns indices of the vector, stably sorted by their value.+--+-- ==== Example+-- >>> import Data.Vector.Algorithms.Intro qualified as VAI+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> argsort $ VU.fromList [0, 1, 0, 1, 0]+-- [0,2,4,1,3]+{-# INLINE argsort #-}+argsort :: (Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int+argsort xs =+  VU.modify+    ( VAI.sortBy+        ( \i j ->+            ( compare (xs VG.! i) (xs VG.! j) <> compare i j+            )+        )+    )+    $ VU.generate (VU.length xs) id++-- | \(O(n)\) Applies a permutation to a mutable vector in-place.+--+-- ==== Constraints+-- - The index array must be a permutation (0-based).+{-# INLINE unsafePermuteInPlace #-}+unsafePermuteInPlace :: (PrimMonad m, VGM.MVector v a) => v (PrimState m) a -> VU.Vector Int -> m ()+unsafePermuteInPlace vec is = stToPrim $ unsafePermuteInPlaceST vec is++-- | \(O(n)\) Applies a permutation to a mutable vector in-place.+--+-- ==== Constraints+-- - The index array must be a permutation (0-based).+{-# INLINEABLE unsafePermuteInPlaceST #-}+unsafePermuteInPlaceST :: (VGM.MVector v a) => v s a -> VU.Vector Int -> ST s ()+unsafePermuteInPlaceST vec is = do+  let !_ = ACIA.runtimeAssert (VGM.length vec == VG.length is) "AtCoder.Extra.Vector.unsafePermuteInPlaceST: the length of the index array must be equal to the length of the permuted vector"+  let inner i lastX = do+        VGM.unsafeWrite vec i lastX+        unless (i == 0) $ do+          let i0' = VG.unsafeIndex is i+          lastX' <- VGM.unsafeRead vec i+          inner i0' lastX'++  let i0' = VG.unsafeIndex is 0+  x0' <- VGM.unsafeRead vec 0+  inner i0' x0'

@@ -99,14 +99,14 @@ -- original array if you can. -- -- @since 1.1.0.0-{-# INLINEABLE access #-}+{-# INLINE access #-} access :: WaveletMatrix -> Int -> Maybe Int access WaveletMatrix {..} i = (xDictWM VG.!) <$> Rwm.access rawWM i  -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\). -- -- @since 1.1.0.0-{-# INLINEABLE rank #-}+{-# INLINE rank #-} rank ::   -- | A wavelet matrix   WaveletMatrix ->@@ -151,7 +151,7 @@ -- not found. -- -- @since 1.1.0.0-{-# INLINEABLE select #-}+{-# INLINE select #-} select :: WaveletMatrix -> Int -> Maybe Int select wm = selectKth wm 0 @@ -180,7 +180,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINEABLE selectIn #-}+{-# INLINE selectIn #-} selectIn ::   -- | A wavelet matrix   WaveletMatrix ->@@ -301,7 +301,7 @@ -- | \(O(\log |S|)\) -- -- @since 1.1.0.0-{-# INLINEABLE unsafeKthSmallestIn #-}+{-# INLINE unsafeKthSmallestIn #-} unsafeKthSmallestIn :: WaveletMatrix -> Int -> Int -> Int -> Int unsafeKthSmallestIn WaveletMatrix {..} l r k =   xDictWM VG.! Rwm.unsafeKthSmallestIn rawWM l r k@@ -334,7 +334,7 @@ -- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0)\). -- -- @since 1.1.0.0-{-# INLINEABLE lookupLT #-}+{-# INLINE lookupLT #-} lookupLT ::   -- | A wavelet matrix   WaveletMatrix ->@@ -376,7 +376,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\). -- -- @since 1.1.0.0-{-# INLINEABLE lookupGT #-}+{-# INLINE lookupGT #-} lookupGT ::   -- | A wavelet matrix   WaveletMatrix ->@@ -394,7 +394,7 @@ -- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINEABLE assocsIn #-}+{-# INLINE assocsIn #-} assocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)] assocsIn WaveletMatrix {..} l r = Rwm.assocsWith rawWM l r (xDictWM VG.!) @@ -402,6 +402,6 @@ -- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINEABLE descAssocsIn #-}+{-# INLINE descAssocsIn #-} descAssocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)] descAssocsIn WaveletMatrix {..} l r = Rwm.descAssocsInWith rawWM l r (xDictWM VG.!)

@@ -249,7 +249,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\). -- -- @since 1.1.0.0-{-# INLINEABLE rank #-}+{-# INLINE rank #-} rank ::   RawWaveletMatrix ->   -- | \(l\)@@ -265,7 +265,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\). -- -- @since 1.1.0.0-{-# INLINEABLE rankBetween #-}+{-# INLINE rankBetween #-} rankBetween ::   RawWaveletMatrix ->   -- | \(l\)@@ -284,7 +284,7 @@ -- not found. -- -- @since 1.1.0.0-{-# INLINEABLE select #-}+{-# INLINE select #-} select :: RawWaveletMatrix -> Int -> Maybe Int select wm = selectKth wm 0 @@ -292,7 +292,7 @@ -- if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINEABLE selectKth #-}+{-# INLINE selectKth #-} selectKth ::   RawWaveletMatrix ->   -- | \(k\)@@ -307,7 +307,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINEABLE selectIn #-}+{-# INLINE selectIn #-} selectIn ::   -- | A wavelet matrix   RawWaveletMatrix ->@@ -453,21 +453,21 @@ -- values are counted as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINEABLE unsafeKthLargestIn #-}+{-# INLINE unsafeKthLargestIn #-} unsafeKthLargestIn :: RawWaveletMatrix -> Int -> Int -> Int -> Int unsafeKthLargestIn wm l r k = unsafeKthSmallestIn wm l r (r - l - (k + 1))  -- | \(O(\log a)\) -- -- @since 1.1.0.0-{-# INLINEABLE unsafeIKthLargestIn #-}+{-# INLINE unsafeIKthLargestIn #-} unsafeIKthLargestIn :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int) unsafeIKthLargestIn wm l r k = unsafeIKthSmallestIn wm l r (r - l - (k + 1))  -- | \(O(\log a)\) -- -- @since 1.1.0.0-{-# INLINEABLE unsafeKthSmallestIn #-}+{-# INLINE unsafeKthSmallestIn #-} unsafeKthSmallestIn :: RawWaveletMatrix -> Int -> Int -> Int -> Int unsafeKthSmallestIn wm l_ r_ k_ =   let (!x, !_, !_, !_) = goDown wm l_ r_ k_@@ -511,7 +511,7 @@ -- | \(O(\log a)\) Finds the maximum \(x\) in \([l, r)\) s.t. \(x_{0} \lt x\). -- -- @since 1.1.0.0-{-# INLINEABLE lookupLT #-}+{-# INLINE lookupLT #-} lookupLT ::   RawWaveletMatrix ->   -- | \(l\)@@ -552,7 +552,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\). -- -- @since 1.1.0.0-{-# INLINEABLE lookupGT #-}+{-# INLINE lookupGT #-} lookupGT ::   RawWaveletMatrix ->   -- | \(l\)@@ -569,7 +569,7 @@ -- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINEABLE assocsIn #-}+{-# INLINE assocsIn #-} assocsIn :: RawWaveletMatrix -> Int -> Int -> [(Int, Int)] assocsIn wm l r = assocsWith wm l r id @@ -612,7 +612,7 @@ -- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINEABLE descAssocsIn #-}+{-# INLINE descAssocsIn #-} descAssocsIn :: RawWaveletMatrix -> Int -> Int -> [(Int, Int)] descAssocsIn wm l r = descAssocsInWith wm l r id 

@@ -4,6 +4,10 @@ -- queries. Points cannot be added after construction, but monoid values in each point can be -- modified later. --+-- ==== SegTree2d vs WaveletMatrix2d+-- They basically the same functionalities and performance, however, in @ac-library-hs@, `SegTree2d`+-- has better API and even outperforms @WaveletMatrix2d@.+-- -- ==== __Example__ -- Create a `WaveletMatrix2d` with initial vertex values: --@@ -200,7 +204,7 @@     i_     $ V.zip (Rwm.bitsRwm rawWmWm2d) segTreesWm2d --- | \(O(\log^2 n)\) Returns the monoid product in \([l, r) \times [y_1, y_2)\).+-- | \(O(\log^2 n)\) Returns monoid product \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\). -- -- @since 1.1.0.0 {-# INLINEABLE prod #-}@@ -218,7 +222,7 @@     !_ = ACIA.checkInterval "AtCoder.Extra.WaveletMatrix.SegTree.prod (compressed x)" xl' xr' (VG.length xDict)     !_ = ACIA.checkInterval "AtCoder.Extra.WaveletMatrix.SegTree.prod (compressed y)" yl' yr' (VG.length yDictWm2d) --- | \(O(\log^2 n)\) Returns the monoid product in \([l, r) \times [y_1, y_2)\). Returns `Nothing` for invalid+-- | \(O(\log^2 n)\) Returns the monoid product in \([x_1, x_2) \times [y_1, y_2)\). Returns `Nothing` for invalid -- intervals. -- -- @since 1.1.0.0@@ -237,7 +241,7 @@     yl' = fromMaybe 0 $ bisectR 0 (VG.length yDictWm2d) $ (< yl) . VG.unsafeIndex yDictWm2d     yr' = fromMaybe (VG.length yDictWm2d) $ bisectR 0 (VG.length yDictWm2d) $ (< yr) . VG.unsafeIndex yDictWm2d --- | \(O(\log^2 n)\) Return the monoid product of all of the points in the wavelet matrix.+-- | \(O(\log^2 n)\) Return the monoid product in \([-\infty, \infty) \times [-\infty, \infty)\). -- -- @since 1.1.0.0 {-# INLINEABLE allProd #-}

@@ -116,13 +116,6 @@ add :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> a -> m () add ft p0 x = stToPrim $ addST ft p0 x --- | \(O(\log n)\) Calculates the sum in a half-open interval @[0, r)@.------ @since 1.0.0.0-{-# INLINE prefixSum #-}-prefixSum :: (PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> m a-prefixSum ft r = stToPrim $ prefixSumST ft r- -- | Calculates the sum in a half-open interval \([l, r)\). -- -- ==== Constraints

@@ -40,6 +40,16 @@ -- *** Exception: AtCoder.Internal.Assert.doctest: given invalid interval `[-1, 0)` over bounds `[0, 5)` -- ... --+-- >>> let !_ = checkPoint2d "AtCoder.Internal.Assert.doctest"  1 1 2 2+-- >>> let !_ = checkPoint2d "AtCoder.Internal.Assert.doctest" 4 4 2 2+-- *** Exception: AtCoder.Internal.Assert.doctest: given invalid point `(4, 4)` for rectangle `[0, 2) x [0, 2)`+-- ...+--+-- >>> let !_ = checkRect "AtCoder.Internal.Assert.doctest"  1 2 1 2 3 3+-- >>> let !_ = checkRect "AtCoder.Internal.Assert.doctest" 1 2 1 2 1 1+-- *** Exception: AtCoder.Internal.Assert.doctest: given invalid rectangle `[1, 2) x [1, 2)` for rectangle `[0, 1) x [0, 1)`+-- ...+-- -- @since 1.0.0.0 module AtCoder.Internal.Assert   ( -- * Runtime assertion@@ -49,6 +59,9 @@     testIndex,     testInterval,     testIntervalBounded,+    testPoint2d,+    testRect,+    testRectShape,      -- * Index assertions     checkIndex,@@ -67,6 +80,14 @@     errorInterval,     checkIntervalBounded,     errorIntervalBounded,++    -- * Two-dimensional index assertions+    checkPoint2d,+    errorPoint2d,+    checkRect,+    errorRect,+    checkRectShape,+    errorRectShape,   ) where @@ -102,6 +123,27 @@ testIntervalBounded :: Int -> Int -> Int -> Int -> Bool testIntervalBounded l r l0 r0 = l0 <= l && l <= r && r <= r0 +-- | \(O(1)\) Tests \((x, y) \in [0, w) \times [0, h)\).+--+-- @since 1.2.3.0+{-# INLINE testPoint2d #-}+testPoint2d :: (HasCallStack) => Int -> Int -> Int -> Int -> Bool+testPoint2d x y w h = 0 <= x && x < w && 0 <= y && y < h++-- | \(O(1)\) Tests \([x_1, x_2) \times [y_1 y_2) \in [0, w) \times [0, h)\).+--+-- @since 1.2.3.0+{-# INLINE testRect #-}+testRect :: (HasCallStack) => Int -> Int -> Int -> Int -> Int -> Int -> Bool+testRect x1 x2 y1 y2 w h = 0 <= x1 && x1 <= x2 && x2 <= w && 0 <= y1 && y1 <= y2 && y2 <= h++-- | \(O(1)\) Tests \(x_1 \le x_2\) and \(y_1 \le \y_2\).+--+-- @since 1.2.3.0+{-# INLINE testRectShape #-}+testRectShape :: (HasCallStack) => Int -> Int -> Int -> Int -> Bool+testRectShape x1 x2 y1 y2 = x1 <= x2 && y1 <= y2+ -- | \(O(1)\) Asserts \(0 \leq i \lt n\) for an array index \(i\). -- -- @since 1.0.0.0@@ -145,7 +187,7 @@   | 0 <= i && i < n = ()   | otherwise = errorVertex funcName i n --- | \(O(1)\) Asserts \(0 \leq i \lt n\) for a graph vertex \(i\).+-- | \(O(1)\) Emits vertex boundary error. -- -- @since 1.0.0.0 {-# INLINE errorVertex #-}@@ -162,7 +204,7 @@   | 0 <= i && i < n = ()   | otherwise = errorEdge funcName i n --- | \(O(1)\) Asserts \(0 \leq i \lt m\) for an edge index \(i\).+-- | \(O(1)\) Emits edge index boundary error. -- -- @since 1.0.0.0 {-# INLINE errorEdge #-}@@ -170,7 +212,7 @@ errorEdge funcName i n =   error $ funcName ++ ": given invalid edge index `" ++ show i ++ "` over the number of edges `" ++ show n ++ "`" --- | \(O(1)\) Asserts \(0 \leq i \lt m\) for an edge index \(i\).+-- | \(O(1)\) Asserts index boundary with custom message. -- -- @since 1.0.0.0 {-# INLINE checkCustom #-}@@ -179,7 +221,7 @@   | testIndex i n = ()   | otherwise = errorCustom funcName indexName i setName n --- | \(O(1)\) Asserts \(0 \leq i \lt m\) for an edge index \(i\).+-- | \(O(1)\) Emis custom index error. -- -- @since 1.0.0.0 {-# INLINE errorCustom #-}@@ -195,7 +237,7 @@   | testInterval l r n = ()   | otherwise = errorInterval funcName l r n --- | \(O(1)\) Asserts \(0 \leq l \leq r \leq n\) for a half-open interval \([l, r)\).+-- | \(O(1)\) Emits interval boundary error. -- -- @since 1.0.0.0 {-# INLINE errorInterval #-}@@ -211,9 +253,60 @@   | testIntervalBounded l r l0 r0 = ()   | otherwise = errorIntervalBounded funcName l r l0 r0 --- | \(O(1)\) Asserts \(0 \leq l \leq r \leq n\) for a half-open interval \([l, r)\).+-- | \(O(1)\) Emits interval boundary error. -- -- @since 1.2.1.0 {-# INLINE errorIntervalBounded #-} errorIntervalBounded :: (HasCallStack) => String -> Int -> Int -> Int -> Int -> a errorIntervalBounded funcName l r l0 r0 = error $ funcName ++ ": given invalid interval `[" ++ show l ++ ", " ++ show r ++ ")` over bounds `[" ++ show l0 ++ ", " ++ show r0 ++ ")`"++-- | \(O(1)\) Asserts \(0 \leq i \lt n\) for a graph vertex \(i\).+--+-- @since 1.2.3.0+{-# INLINE checkPoint2d #-}+checkPoint2d :: (HasCallStack) => String -> Int -> Int -> Int -> Int -> ()+checkPoint2d funcName x y w h+  | testPoint2d x y w h = ()+  | otherwise = errorPoint2d funcName x y w h++-- | \(O(1)\) Emits point boundary error.+--+-- @since 1.2.3.0+{-# INLINE errorPoint2d #-}+errorPoint2d :: (HasCallStack) => String -> Int -> Int -> Int -> Int -> a+errorPoint2d funcName x y w h =+  error $ funcName ++ ": given invalid point `(" ++ show x ++ ", " ++ show y ++ ")` for rectangle `[0, " ++ show w ++ ") x [0, " ++ show h ++ ")`"++-- | \(O(1)\) Asserts \([x_1, x_2) \times [y_1 y_2) \in [0, w) \times [0, h)\).+--+-- @since 1.2.3.0+{-# INLINE checkRect #-}+checkRect :: (HasCallStack) => String -> Int -> Int -> Int -> Int -> Int -> Int -> ()+checkRect funcName x1 x2 y1 y2 w h+  | testRect x1 x2 y1 y2 w h = ()+  | otherwise = errorRect funcName x1 x2 y1 y2 w h++-- | \(O(1)\) Asserts rectangle boundary error.+--+-- @since 1.2.3.0+{-# INLINE errorRect #-}+errorRect :: (HasCallStack) => String -> Int -> Int -> Int -> Int -> Int -> Int -> a+errorRect funcName x1 x2 y1 y2 w h =+  error $ funcName ++ ": given invalid rectangle `[" ++ show x1 ++ ", " ++ show x2 ++ ") x [" ++ show y1 ++ ", " ++ show y2 ++ ")` for rectangle `[0, " ++ show w ++ ") x [0, " ++ show h ++ ")`"++-- | \(O(1)\) Asserts \(x_1 \le x_2\) and \(y_1 \le \y_2\).+--+-- @since 1.2.3.0+{-# INLINE checkRectShape #-}+checkRectShape :: (HasCallStack) => String -> Int -> Int -> Int -> Int -> ()+checkRectShape funcName x1 x2 y1 y2+  | testRectShape x1 x2 y1 y2 = ()+  | otherwise = errorRectShape funcName x1 x2 y1 y2++-- | \(O(1)\) Asserts rectangle boundary error.+--+-- @since 1.2.3.0+{-# INLINE errorRectShape #-}+errorRectShape :: (HasCallStack) => String -> Int -> Int -> Int -> Int -> a+errorRectShape funcName x1 x2 y1 y2 =+  error $ funcName ++ ": given invalid rectangle `[" ++ show x1 ++ ", " ++ show x2 ++ ") x [" ++ show y1 ++ ", " ++ show y2 ++ ")`"

@@ -20,20 +20,101 @@ -- >>> Q.freeze que     -- [_  1, 2  _] -- [1,2] ----- >>> Q.pushFront que 10   -- [10, 1, 2  _]+-- >>> Q.readFront que 0+-- 1+--+-- >>> Q.readFront que 1+-- 2+--+-- >>> Q.readFront que (-1)+-- *** Exception: AtCoder.Internal.Queue.readFront: index out of bounds+-- ...+--+-- >>> Q.readFront que 2+-- *** Exception: AtCoder.Internal.Queue.readFront: index out of bounds+-- ...+--+-- >>> Q.readMaybeFront que (-1)+-- Nothing+--+-- >>> Q.readMaybeFront que 2+-- Nothing+--+-- >>> Q.writeFront que 0 10 -- [_ 10, 2  _]+-- >>> Q.writeFront que 1 20 -- [_ 10, 20 _]+--+-- >>> Q.writeFront que (-1) 777+-- *** Exception: AtCoder.Internal.Queue.modifyFrontM: index out of bounds+-- ...+--+-- >>> Q.writeFront que 2 777+-- *** Exception: AtCoder.Internal.Queue.modifyFrontM: index out of bounds+-- ...+--+-- >>> Q.readBack que 0+-- 20+--+-- >>> Q.readBack que 1+-- 10+--+-- >>> Q.readBack que (-1)+-- *** Exception: AtCoder.Internal.Queue.readBack: index out of bounds+-- ...+--+-- >>> Q.readBack que 2+-- *** Exception: AtCoder.Internal.Queue.readBack: index out of bounds+-- ...+--+-- >>> Q.readMaybeBack que (-1)+-- Nothing+--+-- >>> Q.readMaybeBack que 2+-- Nothing+--+-- >>> Q.writeBack que 0 200+-- >>> Q.writeBack que 1 100+--+-- >>> Q.writeBack que (-1) 777+-- *** Exception: AtCoder.Internal.Queue.modifyBackM: index out of bounds+-- ...+--+-- >>> Q.writeBack que 2 777+-- *** Exception: AtCoder.Internal.Queue.modifyBackM: index out of bounds+-- ...+--+-- >>> Q.pushFront que 10 -- [10, 100, 200  _] -- >>> Q.pushFront que 1000--- *** Exception: AtCoder.Internal.Queue.pushFront: no empty front space+-- *** Exception: AtCoder.Internal.Queue.pushFrontST: no empty front space -- ... ----- >>> Q.unsafeFreeze que -- [10, 1, 2  _]--- [10,1,2]+-- >>> Q.unsafeFreeze que -- [10, 100, 200  _]+-- [10,100,200] ----- >>> Q.clear que      -- [_  _  _  _]+-- >>> Q.clear que        -- [_  _  _  _]+-- >>> Q.peekBack que+-- Nothing+--+-- >>> Q.popFront que+-- Nothing+--+-- >>> Q.popBack que+-- Nothing+-- -- >>> Q.pushBack que 0 -- [0  _  _  _]+-- >>> Q.peekBack que+-- Just 0+-- -- >>> Q.pushBack que 1 -- [0, 1  _  _] -- >>> Q.pushBack que 2 -- [0, 1, 2  _]+-- >>> Q.popBack que    -- [0, 1  _  _]+-- Just 2+-- -- >>> Q.freeze que--- [0,1,2]+-- [0,1]++-- >>> Q.clear que+-- >>> Q.freeze que+-- [] -- -- @since 1.0.0.0 module AtCoder.Internal.Queue@@ -48,15 +129,35 @@     length,     null, -    -- * Modifications+    -- * Element access -    -- ** Push/pop+    -- ** Peek+    peekBack,+    peekFront,++    -- ** Push     pushBack,     pushFront,++    -- ** op+    popBack,+    popBack_,     popFront,     popFront_, -    -- ** Reset+    -- ** Read/write/modify+    readFront,+    readBack,+    readMaybeFront,+    readMaybeBack,+    writeFront,+    writeBack,+    modifyFront,+    modifyFrontM,+    modifyBack,+    modifyBackM,++    -- ** Clear (reset)     clear,      -- * Conversions@@ -65,8 +166,10 @@   ) where +import AtCoder.Internal.Assert qualified as ACIA import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) import Control.Monad.ST (ST)+import Data.Maybe (fromMaybe) import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM@@ -110,6 +213,20 @@ null :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m Bool null = (<$>) (== 0) . length +-- | \(O(1)\) Peeks the last element in the queue.+--+-- @since 1.2.3.0+{-# INLINE peekBack #-}+peekBack :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m (Maybe a)+peekBack que = stToPrim $ peekBackST que++-- | \(O(1)\) Peeks the first element in the queue.+--+-- @since 1.2.3.0+{-# INLINE peekFront #-}+peekFront :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m (Maybe a)+peekFront que = stToPrim $ peekFrontST que+ -- | \(O(1)\) Appends an element to the back. Will throw an exception if the index is out of range. -- -- @since 1.0.0.0@@ -124,6 +241,20 @@ pushFront :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> a -> m () pushFront que e = stToPrim $ pushFrontST que e +-- | \(O(1)\) Removes the last element from the queue and returns it, or `Nothing` if it is empty.+--+-- @since 1.2.3.0+{-# INLINE popBack #-}+popBack :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m (Maybe a)+popBack que = stToPrim $ popBackST que++-- | \(O(1)\) `popBack` with the return value discarded.+--+-- @since 1.2.3.0+{-# INLINE popBack_ #-}+popBack_ :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m ()+popBack_ que = stToPrim $ popBackST_ que+ -- | \(O(1)\) Removes the first element from the queue and returns it, or `Nothing` if it is empty. -- -- @since 1.0.0.0@@ -138,6 +269,92 @@ popFront_ :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m () popFront_ que = stToPrim $ popFrontST_ que +-- | \(O(1)\) Returns the \(k\)-th value from the first element.+--+-- @since 1.2.3.0+{-# INLINE readFront #-}+readFront :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> Int -> m a+readFront que i = stToPrim $ fromMaybe (error msg) <$> readMaybeFrontST que i+  where+    msg = "AtCoder.Internal.Queue.readFront: index out of bounds"++-- | \(O(1)\) Returns the \(k\)-th value from the last element.+--+-- @since 1.2.3.0+{-# INLINE readBack #-}+readBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> Int -> m a+readBack que i = stToPrim $ fromMaybe (error msg) <$> readMaybeBackST que i+  where+    msg = "AtCoder.Internal.Queue.readBack: index out of bounds"++-- | \(O(1)\) Returns the \(k\)-th value from the first element.+--+-- @since 1.2.3.0+{-# INLINE readMaybeFront #-}+readMaybeFront :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> Int -> m (Maybe a)+readMaybeFront que i = stToPrim $ readMaybeFrontST que i++-- | \(O(1)\) Returns the \(k\)-th value from the last element.+--+-- @since 1.2.3.0+{-# INLINE readMaybeBack #-}+readMaybeBack :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> Int -> m (Maybe a)+readMaybeBack que i = stToPrim $ readMaybeBackST que i++-- | \(O(1)\) Writes to the \(k\)-th value from the first element.+--+-- @since 1.2.3.0+{-# INLINE writeFront #-}+writeFront :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> Int -> a -> m ()+writeFront que i x = stToPrim $ do+  modifyFrontM que (pure . const x) i++-- | \(O(1)\) Writes to the \(k\)-th value from the last element.+--+-- @since 1.2.3.0+{-# INLINE writeBack #-}+writeBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> Int -> a -> m ()+writeBack que i x = stToPrim $ do+  modifyBackM que (pure . const x) i++-- | \(O(1)\) Given user function \(f\), returns the \(k\)-th value from the first element.+--+-- @since 1.2.3.0+{-# INLINE modifyFront #-}+modifyFront :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> (a -> a) -> Int -> m ()+modifyFront que f i = stToPrim $ do+  modifyFrontM que (pure . f) i++-- | \(O(1)\) Given user function \(f\), returns the \(k\)-th value from the last element.+--+-- @since 1.2.3.0+{-# INLINE modifyBack #-}+modifyBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> (a -> a) -> Int -> m ()+modifyBack que f i = stToPrim $ do+  modifyBackM que (pure . f) i++-- | \(O(1)\) Given user function \(f\), returns the \(k\)-th value from the first element.+--+-- @since 1.2.3.0+{-# INLINE modifyFrontM #-}+modifyFrontM :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> (a -> m a) -> Int -> m ()+modifyFrontM Queue {..} f i = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  let !_ = ACIA.runtimeAssert (0 <= i && i < r - l) "AtCoder.Internal.Queue.modifyFrontM: index out of bounds"+  VGM.modifyM vecQ f (l + i)++-- | \(O(1)\) Given user function \(f\), returns the \(k\)-th value from the last element.+--+-- @since 1.2.3.0+{-# INLINE modifyBackM #-}+modifyBackM :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> (a -> m a) -> Int -> m ()+modifyBackM Queue {..} f i = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  let !_ = ACIA.runtimeAssert (0 <= i && i < r - l) "AtCoder.Internal.Queue.modifyBackM: index out of bounds"+  VGM.modifyM vecQ f (r - 1 - i)+ -- | \(O(1)\) Sets the `length` to zero. -- -- @since 1.0.0.0@@ -179,6 +396,24 @@   r <- VGM.unsafeRead posQ 1   pure $ r - l +{-# INLINEABLE peekBackST #-}+peekBackST :: (VU.Unbox a) => Queue s a -> ST s (Maybe a)+peekBackST Queue {..} = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  if l >= r+    then pure Nothing+    else Just <$> VGM.read vecQ (r - 1)++{-# INLINEABLE peekFrontST #-}+peekFrontST :: (VU.Unbox a) => Queue s a -> ST s (Maybe a)+peekFrontST Queue {..} = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  if l >= r+    then pure Nothing+    else Just <$> VGM.read vecQ l+ {-# INLINEABLE pushBackST #-} pushBackST :: (HasCallStack, VU.Unbox a) => Queue s a -> a -> ST s () pushBackST Queue {..} e = do@@ -194,17 +429,33 @@ pushFrontST :: (HasCallStack, VU.Unbox a) => Queue s a -> a -> ST s () pushFrontST Queue {..} e = do   l0 <- VGM.unsafeRead posQ 0-  if l0 == 0-    then error "AtCoder.Internal.Queue.pushFront: no empty front space"+  let !_ = ACIA.runtimeAssert (l0 > 0) "AtCoder.Internal.Queue.pushFrontST: no empty front space"+  VGM.unsafeModifyM+    posQ+    ( \l -> do+        VGM.write vecQ (l - 1) e+        pure $ l - 1+    )+    0++{-# INLINEABLE popBackST #-}+popBackST :: (VU.Unbox a) => Queue s a -> ST s (Maybe a)+popBackST Queue {..} = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  if l >= r+    then pure Nothing     else do-      VGM.unsafeModifyM-        posQ-        ( \l -> do-            VGM.write vecQ (l - 1) e-            pure $ l - 1-        )-        0+      x <- VGM.read vecQ (r - 1)+      VGM.unsafeWrite posQ 1 (r - 1)+      pure $ Just x +{-# INLINEABLE popBackST_ #-}+popBackST_ :: (VU.Unbox a) => Queue s a -> ST s ()+popBackST_ que = do+  _ <- popBackST que+  pure ()+ {-# INLINEABLE popFrontST #-} popFrontST :: (VU.Unbox a) => Queue s a -> ST s (Maybe a) popFrontST Queue {..} = do@@ -220,13 +471,26 @@ {-# INLINEABLE popFrontST_ #-} popFrontST_ :: (VU.Unbox a) => Queue s a -> ST s () popFrontST_ que = do-  _ <- popFront que+  _ <- popFrontST que   pure () -{-# INLINEABLE clearST #-}-clearST :: (VU.Unbox a) => Queue s a -> ST s ()-clearST Queue {..} = do-  VGM.set posQ 0+{-# INLINEABLE readMaybeFrontST #-}+readMaybeFrontST :: (VU.Unbox a) => Queue s a -> Int -> ST s (Maybe a)+readMaybeFrontST Queue {..} i = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  if 0 <= i && i < r - l+    then Just <$> VGM.read vecQ (l + i)+    else pure Nothing++{-# INLINEABLE readMaybeBackST #-}+readMaybeBackST :: (VU.Unbox a) => Queue s a -> Int -> ST s (Maybe a)+readMaybeBackST Queue {..} i = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  if 0 <= i && i < r - l+    then Just <$> VGM.read vecQ (r - 1 - i)+    else pure Nothing  {-# INLINEABLE freezeST #-} freezeST :: (VU.Unbox a) => Queue s a -> ST s (VU.Vector a)

@@ -20,6 +20,8 @@ import Tests.Extra.Math qualified import Tests.Extra.Monoid qualified import Tests.Extra.MultiSet qualified+import Tests.Extra.SegTree2d qualified+import Tests.Extra.SegTree2d.Dense qualified import Tests.Extra.Semigroup.Matrix qualified import Tests.Extra.Semigroup.Permutation qualified import Tests.Extra.Seq qualified@@ -71,6 +73,8 @@             testGroup "Math" Tests.Extra.Math.tests,             testGroup "Monoid" Tests.Extra.Monoid.tests,             testGroup "MultiSet" Tests.Extra.MultiSet.tests,+            testGroup "SegTree2d" Tests.Extra.SegTree2d.tests,+            testGroup "SegTree2d.Dense" Tests.Extra.SegTree2d.Dense.tests,             testGroup "Semigroup.Matrix" Tests.Extra.Semigroup.Matrix.tests,             testGroup "Semigroup.Permutation" Tests.Extra.Semigroup.Permutation.tests,             testGroup "Seq" Tests.Extra.Seq.tests,

@@ -0,0 +1,147 @@+{-# LANGUAGE RecordWildCards #-}++module Tests.Extra.SegTree2d (tests) where++import AtCoder.Extra.SegTree2d qualified as Seg+import Control.Monad.Primitive (PrimMonad, PrimState)+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.HUnit+import Test.Tasty.QuickCheck as QC++data Init = Init+  { capacity :: {-# UNPACK #-} !Int,+    q :: {-# UNPACK #-} !Int,+    debugPoints :: VU.Vector (Int, Int),+    vecM :: !(IO (VUM.MVector RealWorld (Int, Int, Sum Int))),+    segM :: !(IO (Seg.SegTree2d RealWorld (Sum Int)))+  }++instance Show Init where+  show Init {..} = show ("Init", capacity, q, debugPoints)++instance QC.Arbitrary Init where+  arbitrary = do+    QC.Positive n <- QC.arbitrary+    q <- QC.chooseInt (1, 2 * n)+    yxs <- (VU.fromList <$>) $ QC.vectorOf n $ do+      x <- QC.chooseInt (-16, 16)+      y <- QC.chooseInt (-16, 16)+      pure (x, y)+    let vecM = VU.thaw $ VU.map (\(!x, !y) -> (x, y, mempty)) yxs+    pure $ Init n q yxs vecM (Seg.new yxs)++data Query+  = -- Read !Int |+    Write !Int !Int+  | ModifyAdd !Int !Int+  | Prod !(Int, Int, Int, Int)+  | AllProd+  | Count !(Int, Int, Int, Int)+  deriving+    ( Eq,+      Show+    )++genQuery :: Int -> QC.Gen Query+genQuery n = do+  QC.oneof+    [ -- Read <$> i,+      Write <$> i <*> val,+      ModifyAdd <$> val <*> i,+      Prod <$> rect,+      pure AllProd,+      Count <$> rect+    ]+  where+    i = QC.chooseInt (0, n - 1)+    rect = do+      xl <- QC.chooseInt (-18, 18)+      xr <- QC.chooseInt (xl, 18)+      yl <- QC.chooseInt (-18, 18)+      yr <- QC.chooseInt (yl, 18)+      pure (xl, xr, yl, yr)+    val = QC.arbitrary @Int++-- | Arbitrary return type for the `Query` result.+data Result+  = None+  | S !(Sum Int)+  | I !Int+  | MS !(Maybe (Sum Int))+  deriving (Show, Eq)++-- | containers. (referencial implementation)+handleRef :: VUM.MVector RealWorld (Int, Int, Sum Int) -> Query -> IO Result+handleRef vec q = case q of+  -- Read i -> do+  --   S . (\(!_, !_, !w) -> w) <$> VGM.read vec i+  Write i v -> do+    (!x, !y, !_) <- VGM.read vec i+    VGM.write vec i (x, y, Sum v)+    pure None+  ModifyAdd w i -> do+    VGM.modify vec (\(!x, !y, !w0) -> (x, y, w0 + Sum w)) i+    pure None+  Prod (!x1, !x2, !y1, !y2) -> do+    vec' <- VU.unsafeFreeze vec+    pure+      . S+      . VU.sum+      . VU.map (\(!_, !_, !w) -> w)+      $ VU.filter (\(!x, !y, !_) -> x1 <= x && x < x2 && y1 <= y && y < y2) vec'+  AllProd -> do+    vec' <- VU.unsafeFreeze vec+    pure . S . VU.sum $ VU.map (\(!_, !_, !w) -> w) vec'+  Count (!x1, !x2, !y1, !y2) -> do+    vec' <- VU.unsafeFreeze vec+    pure+      . I+      . VU.length+      $ VU.filter (\(!x, !y, !_) -> x1 <= x && x < x2 && y1 <= y && y < y2) vec'++handleAcl :: (PrimMonad m) => Seg.SegTree2d (PrimState m) (Sum Int) -> Query -> m Result+handleAcl seg q = case q of+  -- Read i -> do+  --   S <$> Seg.read seg i+  Write i v -> do+    Seg.write seg i $ Sum v+    pure None+  ModifyAdd w i -> do+    Seg.modify seg (+ Sum w) i+    pure None+  Prod (!x1, !x2, !y1, !y2) -> do+    S <$> Seg.prod seg x1 x2 y1 y2+  AllProd -> do+    S <$> Seg.allProd seg+  Count (!x1, !x2, !y1, !y2) -> do+    I <$> Seg.count seg x1 x2 y1 y2++prop_randomTest :: Init -> QC.Property+prop_randomTest Init {..} = QCM.monadicIO $ do+  seg <- QCM.run segM+  vec <- QCM.run vecM+  qs <- QCM.pick $ QC.vectorOf q (genQuery capacity)+  for_ qs $ \query -> do+    expected <- QCM.run $ handleRef vec query+    actual <- QCM.run $ handleAcl seg query+    QCM.assertWith (expected == actual) $ show (query, expected, actual)++unit_zero :: TestTree+unit_zero = testCase "zero" $ do+  seg <- Seg.build @_ @(Sum Int) VU.empty VU.empty VU.empty+  (@?= mempty) =<< Seg.prod seg 0 10 0 10+  (@?= mempty) =<< Seg.allProd seg+  (@?= 0) =<< Seg.count seg 0 10 0 10++tests :: [TestTree]+tests =+  [ unit_zero,+    QC.testProperty "random test" prop_randomTest+  ]

@@ -0,0 +1,146 @@+{-# LANGUAGE RecordWildCards #-}++module Tests.Extra.SegTree2d.Dense (tests) where++import AtCoder.Extra.SegTree2d.Dense qualified as Seg+import Control.Monad.Primitive (PrimMonad, PrimState)+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.HUnit+import Test.Tasty.QuickCheck as QC++data Init = Init+  { wQ :: {-# UNPACK #-} !Int,+    hQ :: {-# UNPACK #-} !Int,+    q :: {-# UNPACK #-} !Int,+    vecM :: !(IO (VUM.MVector RealWorld (Sum Int))),+    segM :: !(IO (Seg.DenseSegTree2d RealWorld (Sum Int)))+  }++instance Show Init where+  show Init {..} = show ("Init", wQ, hQ, q)++instance QC.Arbitrary Init where+  arbitrary = do+    QC.Positive q <- QC.arbitrary+    w <- QC.chooseInt (1, 15)+    h <- QC.chooseInt (1, 15)+    let vecM = VUM.replicate (w * h) mempty+    pure $ Init w h q vecM (Seg.new w h)++data Query+  = Read !(Int, Int)+  | ReadMaybe !(Int, Int)+  | Write !(Int, Int) !Int+  | ModifyAdd !Int !(Int, Int)+  | Prod !(Int, Int, Int, Int)+  | AllProd+  deriving+    ( Eq,+      Show+    )++genQuery :: Int -> Int -> QC.Gen Query+genQuery w h = do+  QC.oneof+    [ Read <$> xy,+      ReadMaybe <$> xy',+      Write <$> xy <*> val,+      ModifyAdd <$> val <*> xy,+      Prod <$> rect,+      pure AllProd+    ]+  where+    xy = (,) <$> QC.chooseInt (0, w - 1) <*> QC.chooseInt (0, h - 1)+    xy' = (,) <$> QC.chooseInt (-1, w) <*> QC.chooseInt (-1, h)+    rect = do+      xl <- QC.chooseInt (-1, w)+      xr <- QC.chooseInt (xl, w)+      yl <- QC.chooseInt (-1, h)+      yr <- QC.chooseInt (yl, h)+      pure (xl, xr, yl, yr)+    val = QC.arbitrary @Int++-- | Arbitrary return type for the `Query` result.+data Result+  = None+  | S !(Sum Int)+  | I !Int+  | MS !(Maybe (Sum Int))+  deriving (Show, Eq)++-- | containers. (referencial implementation)+handleRef :: Int -> Int -> VUM.MVector RealWorld (Sum Int) -> Query -> IO Result+handleRef w h vec q = case q of+  Read (!x, !y) -> do+    S <$> VGM.read vec (w * y + x)+  ReadMaybe (!x, !y)+    | 0 <= x && x < w && 0 <= y && y < h -> MS . Just <$> VGM.read vec (w * y + x)+    | otherwise -> pure $ MS Nothing+  Write (!x, !y) v -> do+    VGM.write vec (w * y + x) (Sum v)+    pure None+  ModifyAdd dw (!x, !y) -> do+    -- FIXME: why die?+    VGM.modify vec (+ Sum dw) (w * y + x)+    pure None+  Prod (!x1, !x2, !y1, !y2) -> do+    vec' <- VU.unsafeFreeze vec+    pure+      . S+      . VU.sum+      $ VU.ifilter+        ( \i _ ->+            let (!y, !x) = i `divMod` w+             in x1 <= x && x < x2 && y1 <= y && y < y2+        )+        vec'+  AllProd -> do+    vec' <- VU.unsafeFreeze vec+    pure . S $ VU.sum vec'++handleAcl :: (PrimMonad m) => Seg.DenseSegTree2d (PrimState m) (Sum Int) -> Query -> m Result+handleAcl seg q = case q of+  Read (!x, !y) -> do+    S <$> Seg.read seg x y+  ReadMaybe (!x, !y) -> do+    MS <$> Seg.readMaybe seg x y+  Write (!x, !y) v -> do+    Seg.write seg x y $ Sum v+    pure None+  ModifyAdd w (!x, !y) -> do+    Seg.modify seg (+ Sum w) x y+    pure None+  Prod (!x1, !x2, !y1, !y2) -> do+    S <$> Seg.prod seg x1 x2 y1 y2+  AllProd -> do+    S <$> Seg.allProd seg++prop_randomTest :: Init -> QC.Property+prop_randomTest Init {..} = QCM.monadicIO $ do+  seg <- QCM.run segM+  vec <- QCM.run vecM+  qs <- QCM.pick $ QC.vectorOf q (genQuery wQ hQ)+  for_ qs $ \query -> do+    expected <- QCM.run $ handleRef wQ hQ vec query+    actual <- QCM.run $ handleAcl seg query+    QCM.assertWith (expected == actual) $ show (query, expected, actual)++unit_zero :: TestTree+unit_zero = testCase "zero" $ do+  seg <- Seg.new @_ @(Sum Int) 0 0+  (@?= mempty) =<< Seg.prod seg 0 10 0 10+  (@?= mempty) =<< Seg.allProd seg+  pure ()++tests :: [TestTree]+tests =+  [ unit_zero,+    QC.testProperty "random test" prop_randomTest+  ]