-- | Square root decomposition is a technique that divides a sequence of values into around
-- \(\sqrt n\) blocks, aggregating the state information for each block. It allows user to process
-- interval query block by block, typically in \(O(\sqrt n)\) time, where a whole block processing
-- take \(O(1)\) time and partial block processing take \(O(\sqrt n)\) time.
--
-- For simplicity, in this document, we assume that highder order functions applided to an entier
-- block (@readFull@ and @actFull@) work in \(O(1)\) time, and those applied to a part of block work
-- in \(O(\sqrt n)\) time. In total, \(q\) query processing takes \(O(q \sqrt n)\) time. Note that
-- it's a rather large number and often requires performance tuning.
--
-- ==== Lazy propagation
-- Typiaclly, an action to a whole block can be delayed; store the aggregation value for the block,
-- delay the internal sequence update, and restore them when part of the block is accessed. Such
-- lazy propagation should be handled on the user side on partial block access functions
-- (@foldPart@ or @actPart@) are called.
--
-- @since 1.2.5.0
module AtCoder.Extra.SqrtDecomposition
  ( -- | These function signatures try to resemble those for lists.
    forM_,
    foldMapM,
    foldMapWithM,
    foldM,
    foldM_,
  )
where

import AtCoder.Internal.Assert qualified as ACIA
import Control.Monad (when)
import Data.Foldable (for_)
import Data.Vector.Unboxed qualified as VU

-- INLINE all the functions, even if the performance gain is just a little bit, in case it matters.

-- | \(O(\sqrt n)\) Runs user function for each block.
{-# INLINE forM_ #-}
forM_ ::
  (Monad m) =>
  -- | Context: block length.
  Int ->
  -- | Function: @actFull@ function that takes target block index.
  (Int -> m ()) ->
  -- | Function: @actPart@ function that takes target block index, left index and right index.
  (Int -> Int -> Int -> m ()) ->
  -- | Input: \(l\).
  Int ->
  -- | Input: \(r\).
  Int ->
  -- | Unit.
  m ()
forM_ :: forall (m :: * -> *).
Monad m =>
Int
-> (Int -> m ())
-> (Int -> Int -> Int -> m ())
-> Int
-> Int
-> m ()
forM_ !Int
blockLen !Int -> m ()
actFull !Int -> Int -> Int -> m ()
actPart !Int
l !Int
r = do
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r) String
"AtCoder.Extra.SqrtDecomposition.forM_: `l <= r` must hold"
  let (!Int
il, !Int
remL) = Int
l Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
blockLen
  let (!Int
ir, !Int
remR) = Int
r Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
blockLen
  if Int
il Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ir
    then do
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
remR Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
remL) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        Int -> Int -> Int -> m ()
actPart Int
il Int
l Int
r
    else do
      if Int
remL Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then Int -> m ()
actFull Int
il
        else Int -> Int -> Int -> m ()
actPart Int
il Int
l (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
remL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
blockLen)
      [Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
il Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 .. Int
ir Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
iBlock -> do
        Int -> m ()
actFull Int
iBlock
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
remR Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        Int -> Int -> Int -> m ()
actPart Int
ir (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
remR) Int
r

-- | \(O(\sqrt n)\) Runs user function for each block and concatanate their monoid output.
--
-- ==== Constraints
-- - \(l \le r\)
-- - If an empty interval is queried, the @readPart@ function must return a valid value.
--
-- @since 1.2.5.0
{-# INLINE foldMapM #-}
foldMapM ::
  (Monad m, Semigroup a) =>
  -- | Context: block length.
  Int ->
  -- | Function: @readFull@ function that takes target block index and returns monoid value of it.
  (Int -> m a) ->
  -- | Function: @readPart@ function that takes target block index, left index and right index, and
  -- returns monoid value for it.
  (Int -> Int -> Int -> m a) ->
  -- | Input: \(l\).
  Int ->
  -- | Input: \(r\).
  Int ->
  -- | Concatenated output.
  m a
foldMapM :: forall (m :: * -> *) a.
(Monad m, Semigroup a) =>
Int
-> (Int -> m a) -> (Int -> Int -> Int -> m a) -> Int -> Int -> m a
foldMapM Int
blockLen = Int
-> (a -> a -> a)
-> (Int -> m a)
-> (Int -> Int -> Int -> m a)
-> Int
-> Int
-> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> (a -> a -> a)
-> (Int -> m a)
-> (Int -> Int -> Int -> m a)
-> Int
-> Int
-> m a
foldMapWithM Int
blockLen a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>)

-- | \(O(\sqrt n)\) Runs user function for each block and concatanates their output with user
-- function.
--
-- ==== Constraints
-- - \(l \le r\)
-- - If an empty interval is queried, the @readPart@ function must return a valid value.
--
-- @since 1.2.5.0
{-# INLINE foldMapWithM #-}
foldMapWithM ::
  (Monad m) =>
  -- | Context: block length.
  Int ->
  -- | Merges function for output values.
  (a -> a -> a) ->
  -- | Function: @readFull@ function that takes target block index and returns monoid value of it.
  (Int -> m a) ->
  -- | Function: @readPart@ function that takes target block index, left index and right index, and
  -- returns output value of it.
  (Int -> Int -> Int -> m a) ->
  -- | Input: \(l\).
  Int ->
  -- | Input: \(r\).
  Int ->
  -- | Concatenated output.
  m a
foldMapWithM :: forall (m :: * -> *) a.
Monad m =>
Int
-> (a -> a -> a)
-> (Int -> m a)
-> (Int -> Int -> Int -> m a)
-> Int
-> Int
-> m a
foldMapWithM !Int
blockLen !a -> a -> a
merge !Int -> m a
readFull !Int -> Int -> Int -> m a
readPart !Int
l !Int
r = do
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r) String
"AtCoder.Extra.SqrtDecomposition.foldMapWithM: `l <= r` must hold"
  let (!Int
il, !Int
remL) = Int
l Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
blockLen
  let (!Int
ir, !Int
remR) = Int
r Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
blockLen
  if Int
il Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ir
    then do
      Int -> Int -> Int -> m a
readPart Int
il Int
l Int
r
    else do
      !a
sx <-
        if Int
remL Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
          then Int -> m a
readFull Int
il
          else Int -> Int -> Int -> m a
readPart Int
il Int
l (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
remL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
blockLen)
      !a
sm <-
        (a -> Int -> m a) -> a -> Vector Int -> m a
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM'
          (\ !a
acc Int
iBlock -> a -> a -> a
merge a
acc (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m a
readFull Int
iBlock)
          a
sx
          (Vector Int -> m a) -> Vector Int -> m a
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Int
ir Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
il) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
il Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
      if Int
remR Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
sm
        else do
          a
rx <- Int -> Int -> Int -> m a
readPart Int
ir (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
remR) Int
r
          a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$! a -> a -> a
merge a
sm a
rx

-- | \(O(\sqrt n)\) Runs user function for each block, performing left folding.
--
-- ==== Constraints
-- - \(l \le r\)
--
-- @since 1.2.5.0
{-# INLINE foldM #-}
foldM ::
  (Monad m) =>
  -- | Context: block length.
  Int ->
  -- | Function: @foldFull@ function that takes target block index and returns monoid value of it.
  (a -> Int -> m a) ->
  -- | Function: @foldPart@ function that takes target block index, left and right local index and returns monoid
  -- value of it.
  (a -> Int -> Int -> Int -> m a) ->
  -- | Initial folding value.
  a ->
  -- | Input: \(l\).
  Int ->
  -- | Input: \(r\).
  Int ->
  -- | Folding result.
  m a
foldM :: forall (m :: * -> *) a.
Monad m =>
Int
-> (a -> Int -> m a)
-> (a -> Int -> Int -> Int -> m a)
-> a
-> Int
-> Int
-> m a
foldM !Int
blockLen !a -> Int -> m a
foldFull !a -> Int -> Int -> Int -> m a
foldPart !a
s0 !Int
l !Int
r = do
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r) String
"AtCoder.Extra.SqrtDecomposition.foldM: `l <= r` must hold"
  let (!Int
il, !Int
remL) = Int
l Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
blockLen
  let (!Int
ir, !Int
remR) = Int
r Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
blockLen
  if Int
il Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ir
    then do
      if Int
remL Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
remR
        then a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
s0
        else a -> Int -> Int -> Int -> m a
foldPart a
s0 Int
il Int
l Int
r
    else do
      !a
sx <-
        if Int
remL Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
          then a -> Int -> m a
foldFull a
s0 Int
il
          else a -> Int -> Int -> Int -> m a
foldPart a
s0 Int
il Int
l (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
remL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
blockLen)
      !a
sm <-
        (a -> Int -> m a) -> a -> Vector Int -> m a
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM'
          a -> Int -> m a
foldFull
          a
sx
          (Vector Int -> m a) -> Vector Int -> m a
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Int
ir Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
il) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
il Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
      if Int
remR Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
sm
        else a -> Int -> Int -> Int -> m a
foldPart a
sm Int
ir (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
remR) Int
r

-- | \(O(\sqrt n)\) `foldM` with return value discarded.
--
-- ==== Constraints
-- - \(l \le r\)
--
-- @since 1.2.5.0
{-# INLINE foldM_ #-}
foldM_ ::
  (Monad m) =>
  -- | Context: Block length.
  Int ->
  -- | @readFull@ function that takes target block index and returns monoid value of it.
  (a -> Int -> m a) ->
  -- | @readPart@ function that takes target block index, left and right local index and returns monoid
  -- value of it.
  (a -> Int -> Int -> Int -> m a) ->
  -- | Initial folding value.
  a ->
  -- | Input: \(l\).
  Int ->
  -- | Input: \(r\).
  Int ->
  -- | Unit.
  m ()
foldM_ :: forall (m :: * -> *) a.
Monad m =>
Int
-> (a -> Int -> m a)
-> (a -> Int -> Int -> Int -> m a)
-> a
-> Int
-> Int
-> m ()
foldM_ !Int
blockLen !a -> Int -> m a
readFull !a -> Int -> Int -> Int -> m a
readPart !a
s0 !Int
l !Int
r = do
  a
_ <- Int
-> (a -> Int -> m a)
-> (a -> Int -> Int -> Int -> m a)
-> a
-> Int
-> Int
-> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> (a -> Int -> m a)
-> (a -> Int -> Int -> Int -> m a)
-> a
-> Int
-> Int
-> m a
foldM Int
blockLen a -> Int -> m a
readFull a -> Int -> Int -> Int -> m a
readPart a
s0 Int
l Int
r
  () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()