{-# 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
    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 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 = 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 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
  | 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 \(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 :: 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 \(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 :: 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
  -- FIXME: correct?
  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)