{-# 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 Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
import Data.Bits
import Data.Foldable (for_)
import Data.Maybe (fromMaybe)
import Data.Vector qualified as V
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
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
    forall s a. DenseSegTree2d s a -> Int
hDst :: {-# UNPACK #-} !Int,
    -- | Width
    --
    -- @since 1.2.3.0
    forall s a. DenseSegTree2d s a -> Int
wDst :: {-# UNPACK #-} !Int,
    -- | Monoid values
    --
    -- @since 1.2.3.0
    forall s a. DenseSegTree2d s a -> MVector s a
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 :: forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Int -> Int -> m (DenseSegTree2d (PrimState m) a)
new Int
wDst Int
hDst = ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (DenseSegTree2d (PrimState m) a)
 -> m (DenseSegTree2d (PrimState m) a))
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ do
  MVector (PrimState m) a
dataDst <- Int
-> a -> ST (PrimState m) (MVector (PrimState (ST (PrimState m))) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDst) a
forall a. Monoid a => a
mempty
  DenseSegTree2d (PrimState m) a
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
wDst :: Int
hDst :: Int
dataDst :: MVector (PrimState m) a
..}

-- | \(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 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Int -> Int -> Vector a -> m (DenseSegTree2d (PrimState m) a)
build Int
w Int
h Vector a
xs = ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (DenseSegTree2d (PrimState m) a)
 -> m (DenseSegTree2d (PrimState m) a))
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Vector (Vector a)
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Vector (Vector a) -> ST s (DenseSegTree2d s a)
buildST (Vector (Vector a)
 -> ST (PrimState m) (DenseSegTree2d (PrimState m) a))
-> Vector (Vector a)
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Int
-> (Vector a -> (Vector a, Vector a))
-> Vector a
-> Vector (Vector a)
forall b a. Int -> (b -> (a, b)) -> b -> Vector a
V.unfoldrExactN Int
h (Int -> Vector a -> (Vector a, Vector a)
forall a. Unbox a => Int -> Vector a -> (Vector a, Vector a)
VU.splitAt Int
w) Vector a
xs
  where
    !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h) String
"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' :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Vector (Vector a) -> m (DenseSegTree2d (PrimState m) a)
build' Vector (Vector a)
xs = ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (DenseSegTree2d (PrimState m) a)
 -> m (DenseSegTree2d (PrimState m) a))
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Vector (Vector a)
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Vector (Vector a) -> ST s (DenseSegTree2d s a)
buildST Vector (Vector a)
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 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> Int -> Int -> m a
read DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} Int
x Int
y = do
  let !()
_ = HasCallStack => String -> Int -> Int -> Int -> Int -> ()
String -> Int -> Int -> Int -> Int -> ()
ACIA.checkPoint2d String
"AtCoder.Extra.SegTree2d.Dense.read" Int
x Int
y Int
wDst Int
hDst
  MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
dataDst (Int -> m a) -> Int -> m a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Int
idx Int
wDst (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst) (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
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 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> Int -> Int -> m (Maybe a)
readMaybe DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} Int
x Int
y
  | HasCallStack => Int -> Int -> Int -> Int -> Bool
Int -> Int -> Int -> Int -> Bool
ACIA.testPoint2d Int
x Int
y Int
wDst Int
hDst = do
      a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst) (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wDst))
  | Bool
otherwise = Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
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 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> Int -> Int -> a -> m ()
write seg :: DenseSegTree2d (PrimState m) a
seg@DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} Int
x Int
y a
a = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  DenseSegTree2d (PrimState (ST (PrimState m))) a
-> (a -> ST (PrimState m) a) -> Int -> Int -> ST (PrimState m) ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyM DenseSegTree2d (PrimState m) a
DenseSegTree2d (PrimState (ST (PrimState m))) a
seg (a -> ST (PrimState m) a
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST (PrimState m) a) -> (a -> a) -> a -> ST (PrimState m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> a
forall a b. a -> b -> a
const a
a) Int
x Int
y

-- | \(O(\log  h \log w)\) Given a user function \(f\), modifies the monoid value at \((x, y)\) with
-- it.
--
-- @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 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> (a -> a) -> Int -> Int -> m ()
modify DenseSegTree2d (PrimState m) a
seg a -> a
f Int
x Int
y = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  DenseSegTree2d (PrimState (ST (PrimState m))) a
-> (a -> ST (PrimState m) a) -> Int -> Int -> ST (PrimState m) ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyM DenseSegTree2d (PrimState m) a
DenseSegTree2d (PrimState (ST (PrimState m))) a
seg (a -> ST (PrimState m) a
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST (PrimState m) a) -> (a -> a) -> a -> ST (PrimState m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f) Int
x Int
y

-- | \(O(\log h \log w)\) Given a user function \(f\), modifies the monoid value at \((x, y)\) with
-- it.
--
-- @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 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyM DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} a -> m a
f Int
x0_ Int
y0_ = do
  let !()
_ = HasCallStack => String -> Int -> Int -> Int -> Int -> ()
String -> Int -> Int -> Int -> Int -> ()
ACIA.checkPoint2d String
"AtCoder.Extra.SegTree2d.Dense.modifyM" Int
x0_ Int
y0_ Int
wDst Int
hDst
  let y0 :: Int
y0 = Int
y0_ Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst
  let x0 :: Int
x0 = Int
x0_ Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wDst
  MVector (PrimState m) a -> (a -> m a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.modifyM MVector (PrimState m) a
dataDst a -> m a
f (Int -> Int -> Int -> Int
idx Int
wDst Int
y0 Int
x0)
  ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    -- right to left
    let updateCurrentRow :: Int -> f ()
updateCurrentRow Int
0 = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        updateCurrentRow Int
x = do
          a
xl <- MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y0 (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0))
          a
xr <- MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y0 (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
          MVector (PrimState f) a -> Int -> a -> f ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y0 Int
x) (a -> f ()) -> a -> f ()
forall a b. (a -> b) -> a -> b
$! a
xl a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
xr
          Int -> f ()
updateCurrentRow (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
    Int -> ST (PrimState m) ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> f ()
updateCurrentRow (Int
x0 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)

    -- down to up
    let updateOtherRow :: Int -> f ()
updateOtherRow Int
0 = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        updateOtherRow Int
y = do
          let updateRow :: Int -> f ()
updateRow Int
0 = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              updateRow Int
x = do
                a
xl <- MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Int
x)
                a
xr <- MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
x)
                MVector (PrimState f) a -> Int -> a -> f ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
x) (a -> f ()) -> a -> f ()
forall a b. (a -> b) -> a -> b
$! a
xl a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
xr
                Int -> f ()
updateRow (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
          Int -> f ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> f ()
updateRow Int
x0
          Int -> f ()
updateOtherRow (Int
y Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
    Int -> ST (PrimState m) ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> f ()
updateOtherRow (Int
y0 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
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 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a
prod seg :: DenseSegTree2d (PrimState m) a
seg@DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} Int
x1 Int
x2 Int
y1 Int
y2 = ST (PrimState m) a -> m a
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) a -> m a) -> ST (PrimState m) a -> m a
forall a b. (a -> b) -> a -> b
$ do
  let !()
_ = HasCallStack => String -> Int -> Int -> Int -> Int -> ()
String -> Int -> Int -> Int -> Int -> ()
ACIA.checkRectShape String
"AtCoder.Extra.SegTree2d.Dense.prodST" Int
x1 Int
x2 Int
y1 Int
y2
  DenseSegTree2d (PrimState m) a
-> Int -> Int -> Int -> Int -> ST (PrimState m) a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
prodST DenseSegTree2d (PrimState m) a
seg (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
x1) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
wDst Int
x2) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
y1) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
hDst Int
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 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> m a
allProd DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} = ST (PrimState m) a -> m a
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) a -> m a) -> ST (PrimState m) a -> m a
forall a b. (a -> b) -> a -> b
$ do
  a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
forall a. Monoid a => a
mempty (Maybe a -> a) -> ST (PrimState m) (Maybe a) -> ST (PrimState m) a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST (PrimState m))) a
-> Int -> ST (PrimState m) (Maybe a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (Maybe a)
VGM.readMaybe MVector (PrimState m) a
MVector (PrimState (ST (PrimState m))) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
1 Int
1)

-- -------------------------------------------------------------------------------------------------
-- Internal
-- -------------------------------------------------------------------------------------------------

{-# INLINE idx #-}
idx :: Int -> Int -> Int -> Int
idx :: Int -> Int -> Int -> Int
idx Int
w Int
y Int
x = Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
w) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x

{-# INLINEABLE buildST #-}
buildST :: (HasCallStack, Monoid a, VU.Unbox a) => V.Vector (VU.Vector a) -> ST s (DenseSegTree2d s a)
buildST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Vector (Vector a) -> ST s (DenseSegTree2d s a)
buildST Vector (Vector a)
vec = do
  let hDst :: Int
hDst = Vector (Vector a) -> Int
forall a. Vector a -> Int
V.length Vector (Vector a)
vec
  let wDst :: Int
wDst = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length (Vector (Vector a) -> Vector a
forall a. Vector a -> a
V.head Vector (Vector a)
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
  MVector s a
dataDst <- Int -> a -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
wDst) a
forall a. Monoid a => a
mempty

  -- copy the base matrix in [w, 2w) \times [h, 2h):
  Vector (Vector a) -> (Int -> Vector a -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
V.iforM_ Vector (Vector a)
vec ((Int -> Vector a -> ST s ()) -> ST s ())
-> (Int -> Vector a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
y Vector a
vs -> do
    Vector a -> (Int -> a -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector a
vs ((Int -> a -> ST s ()) -> ST s ())
-> (Int -> a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x a
v -> do
      MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y) (Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)) a
v

  -- extend the row (y >= h) as a segment tree's internal vector:
  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
hDst .. Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
y -> do
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
0] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x -> do
      a
xl <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0))
      a
xr <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
      MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
x) (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a
xl a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
xr

  -- extend each column as a segment tree:
  -- NOTE (pref): interate from y then x for contiguous memory access
  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
0] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
y -> do
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x -> do
      a
xl <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Int
x)
      a
xr <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
x)
      MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
x) (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a
xl a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
xr

  DenseSegTree2d s a -> ST s (DenseSegTree2d s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DenseSegTree2d {Int
MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector s a
..}

{-# INLINEABLE prodST #-}
prodST :: (HasCallStack, Monoid a, VU.Unbox a) => DenseSegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
prodST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
prodST seg :: DenseSegTree2d s a
seg@DenseSegTree2d {Int
MVector s a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector s a
..} Int
x1 Int
x2 Int
y1 Int
y2 = do
  a -> Int -> Int -> ST s a
inner a
forall a. Monoid a => a
mempty (Int
y1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst) (Int
y2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    -- inclusive interval [yl, yr]
    inner :: a -> Int -> Int -> ST s a
inner !a
acc Int
yl Int
yr
      | Int
yl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
yr = a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
      | Bool
otherwise = do
          a
acc' <-
            if Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
yl Int
0
              then (a
acc <>) (a -> a) -> ST s a -> ST s a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
prodX DenseSegTree2d s a
seg Int
yl Int
x1 Int
x2
              else a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
          a
acc'' <-
            if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
yr Int
0
              then (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
acc') (a -> a) -> ST s a -> ST s a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
prodX DenseSegTree2d s a
seg Int
yr Int
x1 Int
x2
              else a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc'
          a -> Int -> Int -> ST s a
inner a
acc'' ((Int
yl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
1) ((Int
yr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
1)

{-# INLINEABLE prodX #-}
prodX :: (HasCallStack, Monoid a, VU.Unbox a) => DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
prodX :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
prodX DenseSegTree2d {Int
MVector s a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector s a
..} Int
y Int
x1 Int
x2 = do
  a -> Int -> Int -> ST s a
forall {f :: * -> *}.
(PrimState f ~ s, PrimMonad f) =>
a -> Int -> Int -> f a
inner a
forall a. Monoid a => a
mempty (Int
x1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wDst) (Int
x2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    -- inclusive interval [xl, xr]
    inner :: a -> Int -> Int -> f a
inner !a
acc Int
xl Int
xr
      | Int
xl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
xr = a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
      | Bool
otherwise =do
          a
acc' <-
            if Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
xl Int
0
              then (a
acc <>) (a -> a) -> f a -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
xl)
              else a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
          a
acc'' <-
            if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
xr Int
0
              then (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
acc') (a -> a) -> f a -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
xr)
              else a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc'
          a -> Int -> Int -> f a
inner a
acc'' ((Int
xl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
1) ((Int
xr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
1)