{-# LANGUAGE RecordWildCards #-}

-- | A Fenwick tree, also known as binary indexed tree. Given an array of length \(n\), it processes
-- the following queries in \(O(\log n)\) time.
--
-- - Updating an element
-- - Calculating the sum of the elements of an interval
--
-- ==== __Example__
-- Create a `FenwickTree` with `new`:
--
-- >>> import AtCoder.FenwickTree qualified as FT
-- >>> ft <- FT.new @_ @Int 4 -- [0, 0, 0, 0]
-- >>> FT.nFt ft
-- 4
--
-- It can perform point `add` and range `sum` in \(O(\log n)\) time:
--
-- >>> FT.add ft 0 3          -- [3, 0, 0, 0]
-- >>> FT.sum ft 0 3
-- 3
--
-- >>> FT.add ft 2 3          -- [3, 0, 3, 0]
-- >>> FT.sum ft 0 3
-- 6
--
-- Create a `FenwickTree` with initial values using `build`:
--
-- >>> ft <- FT.build @_ @Int $ VU.fromList [3, 0, 3, 0]
-- >>> FT.add ft 1 2          -- [3, 2, 3, 0]
-- >>> FT.sum ft 0 3
-- 8
--
-- @since 1.0.0.0
module AtCoder.FenwickTree
  ( -- * Fenwick tree
    FenwickTree (nFt),

    -- * Constructors
    new,
    build,

    -- * Adding
    add,

    -- * Accessors
    sum,
    sumMaybe,

    -- * Bisection methods
    maxRight,
    maxRightM,
    minLeft,
    minLeftM,
  )
where

import AtCoder.Internal.Assert qualified as ACIA
import Control.Monad (when)
import Control.Monad.Fix (fix)
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Data.Bits
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 (sum)

-- | A Fenwick tree.
--
-- @since 1.0.0.0
data FenwickTree s a = FenwickTree
  { -- | 1.0.0 The number of vertices.
    --
    -- @since 1.0.0.0
    forall s a. FenwickTree s a -> Int
nFt :: {-# UNPACK #-} !Int,
    -- | The data storage.
    forall s a. FenwickTree s a -> MVector s a
dataFt :: !(VUM.MVector s a)
  }

-- | Creates an array \([a_0, a_1, \cdots, a_{n-1}]\) of length \(n\). All the elements are
-- initialized to \(0\).
--
-- ==== Constraints
-- - \(0 \leq n\)
--
-- ==== Complexity
-- - \(O(n)\)
--
-- @since 1.0.0.0
{-# INLINE new #-}
new :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => Int -> m (FenwickTree (PrimState m) a)
new :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
Int -> m (FenwickTree (PrimState m) a)
new Int
nFt
  | Int
nFt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 = do
      MVector (PrimState m) a
dataFt <- Int -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nFt a
0
      FenwickTree (PrimState m) a -> m (FenwickTree (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FenwickTree {Int
MVector (PrimState m) a
nFt :: Int
dataFt :: MVector (PrimState m) a
nFt :: Int
dataFt :: MVector (PrimState m) a
..}
  | Bool
otherwise = [Char] -> m (FenwickTree (PrimState m) a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> m (FenwickTree (PrimState m) a))
-> [Char] -> m (FenwickTree (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ [Char]
"AtCoder.FenwickTree.new: given negative size `" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
nFt [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"`"

-- | Creates `FenwickTree` with initial values.
--
-- ==== Complexity
-- - \(O(n)\)
--
-- @since 1.0.0.0
build :: (PrimMonad m, Num a, VU.Unbox a) => VU.Vector a -> m (FenwickTree (PrimState m) a)
{-# INLINE build #-}
build :: forall (m :: * -> *) a.
(PrimMonad m, Num a, Unbox a) =>
Vector a -> m (FenwickTree (PrimState m) a)
build Vector a
xs = do
  FenwickTree (PrimState m) a
ft <- Int -> m (FenwickTree (PrimState m) a)
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
Int -> m (FenwickTree (PrimState m) a)
new (Int -> m (FenwickTree (PrimState m) a))
-> Int -> m (FenwickTree (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs
  Vector a -> (Int -> a -> m ()) -> m ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector a
xs ((Int -> a -> m ()) -> m ()) -> (Int -> a -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ FenwickTree (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> a -> m ()
add FenwickTree (PrimState m) a
ft
  FenwickTree (PrimState m) a -> m (FenwickTree (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FenwickTree (PrimState m) a
ft

-- | Adds \(x\) to \(p\)-th value of the array.
--
-- ==== Constraints
-- - \(0 \leq l \lt n\)
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- @since 1.0.0.0
{-# INLINE add #-}
add :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> a -> m ()
add :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> a -> m ()
add FenwickTree {Int
MVector (PrimState m) a
nFt :: forall s a. FenwickTree s a -> Int
dataFt :: forall s a. FenwickTree s a -> MVector s a
nFt :: Int
dataFt :: MVector (PrimState m) a
..} Int
p0 a
x = do
  let !()
_ = HasCallStack => [Char] -> Int -> Int -> ()
[Char] -> Int -> Int -> ()
ACIA.checkIndex [Char]
"AtCoder.FenwickTree.add" Int
p0 Int
nFt
  let p1 :: Int
p1 = Int
p0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  (((Int -> m ()) -> Int -> m ()) -> Int -> m ())
-> Int -> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> m ()) -> Int -> m ()) -> Int -> m ()
forall a. (a -> a) -> a
fix Int
p1 (((Int -> m ()) -> Int -> m ()) -> m ())
-> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int -> m ()
loop Int
p -> do
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
nFt) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      MVector (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState m) a
dataFt (a -> a -> a
forall a. Num a => a -> a -> a
+ a
x) (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
      Int -> m ()
loop (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$! Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
p Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
p))

-- | \(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 :: forall (m :: * -> *) a.
(PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> m a
prefixSum FenwickTree {Int
MVector (PrimState m) a
nFt :: forall s a. FenwickTree s a -> Int
dataFt :: forall s a. FenwickTree s a -> MVector s a
nFt :: Int
dataFt :: MVector (PrimState m) a
..} = a -> Int -> m a
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
a -> Int -> f a
inner a
0
  where
    inner :: a -> Int -> f a
inner !a
acc !Int
r
      | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
      | Bool
otherwise = do
          a
dx <- 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
dataFt (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          a -> Int -> f a
inner (a
acc a -> a -> a
forall a. Num a => a -> a -> a
+ a
dx) (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
r))

-- | Calculates the sum in a half-open interval \([l, r)\).
--
-- ==== Constraints
-- - \(0 \leq l \leq r \leq n\)
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- @since 1.0.0.0
{-# INLINE sum #-}
sum :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> Int -> m a
sum :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> Int -> m a
sum ft :: FenwickTree (PrimState m) a
ft@FenwickTree {Int
nFt :: forall s a. FenwickTree s a -> Int
nFt :: Int
nFt} Int
l Int
r
  | Bool -> Bool
not (Int -> Int -> Int -> Bool
ACIA.testInterval Int
l Int
r Int
nFt) = [Char] -> Int -> Int -> Int -> m a
forall a. HasCallStack => [Char] -> Int -> Int -> Int -> a
ACIA.errorInterval [Char]
"AtCoder.FenwickTree.sum" Int
l Int
r Int
nFt
  | Bool
otherwise = FenwickTree (PrimState m) a -> Int -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> Int -> m a
unsafeSum FenwickTree (PrimState m) a
ft Int
l Int
r

-- | Total variant of `sum`. Calculates the sum in a half-open interval \([l, r)\). It returns
-- `Nothing` if the interval is invalid.
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- @since 1.0.0.0
{-# INLINE sumMaybe #-}
sumMaybe :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> Int -> m (Maybe a)
sumMaybe :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> Int -> m (Maybe a)
sumMaybe ft :: FenwickTree (PrimState m) a
ft@FenwickTree {Int
nFt :: forall s a. FenwickTree s a -> Int
nFt :: Int
nFt} Int
l Int
r
  | Bool -> Bool
not (Int -> Int -> Int -> Bool
ACIA.testInterval Int
l Int
r Int
nFt) = 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
  | Bool
otherwise = 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
<$> FenwickTree (PrimState m) a -> Int -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> Int -> m a
unsafeSum FenwickTree (PrimState m) a
ft Int
l Int
r

-- | Internal implementation of `sum`.
{-# INLINE unsafeSum #-}
unsafeSum :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> Int -> m a
unsafeSum :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> Int -> m a
unsafeSum FenwickTree (PrimState m) a
ft Int
l Int
r = do
  a
xr <- FenwickTree (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> m a
prefixSum FenwickTree (PrimState m) a
ft Int
r
  a
xl <- FenwickTree (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> m a
prefixSum FenwickTree (PrimState m) a
ft Int
l
  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
xr a -> a -> a
forall a. Num a => a -> a -> a
- a
xl

-- | (Extra API) Applies a binary search on the Fenwick tree. It returns an index \(r\) that
-- satisfies both of the following.
--
-- - \(r = l\) or \(f(a[l] + a[l + 1] + ... + a[r - 1])\) returns `True`.
-- - \(r = n\) or \(f(a[l] + a[l + 1] + ... + a[r]))\) returns `False`.
--
-- If \(f\) is monotone, this is the maximum \(r\) that satisfies
-- \(f(a[l] + a[l + 1] + ... + a[r - 1])\).
--
-- ==== Constraints
-- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side effect.
-- - \(f(0)\) returns `True`
-- - \(0 \leq l \leq n\)
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- @since 1.2.2.0
{-# INLINE maxRight #-}
maxRight ::
  (HasCallStack, PrimMonad m, Num a, VU.Unbox a) =>
  -- | The Fenwick tree
  FenwickTree (PrimState m) a ->
  -- | \(l\)
  Int ->
  -- | \(p\): user predicate
  (a -> Bool) ->
  -- | \(r\): \(p\) holds for \([l, r)\)
  m Int
maxRight :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> (a -> Bool) -> m Int
maxRight FenwickTree (PrimState m) a
ft Int
l0 a -> Bool
f = FenwickTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
maxRightM FenwickTree (PrimState m) a
ft Int
l0 (Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> m Bool) -> (a -> Bool) -> a -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
f)

-- | (Extra API) Monadic variant of `maxRight`.
--
-- ==== Constraints
-- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side effect.
-- - \(f(0)\) returns `True`
-- - \(0 \leq l \leq n\)
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- @since 1.2.2.0
{-# INLINEABLE maxRightM #-}
maxRightM ::
  forall m a.
  (HasCallStack, PrimMonad m, Num a, VU.Unbox a) =>
  -- | The Fenwick tree
  FenwickTree (PrimState m) a ->
  -- | \(l\)
  Int ->
  -- | \(p\): user predicate
  (a -> m Bool) ->
  -- | \(r\): \(p\) holds for \([l, r)\)
  m Int
maxRightM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
maxRightM FenwickTree {Int
MVector (PrimState m) a
nFt :: forall s a. FenwickTree s a -> Int
dataFt :: forall s a. FenwickTree s a -> MVector s a
nFt :: Int
dataFt :: MVector (PrimState m) a
..} Int
l0 a -> m Bool
f = do
  Bool
b0 <- a -> m Bool
f a
0
  let !()
_ = HasCallStack => Bool -> [Char] -> ()
Bool -> [Char] -> ()
ACIA.runtimeAssert Bool
b0 [Char]
"AtCoder.FenwickTree.maxRightM: `f 0` must return `True`"

  let inner :: Int -> a -> m (Int, Int, a)
inner Int
i !a
s
        | Int -> Bool
forall a. Integral a => a -> Bool
odd Int
i = do
            a
ds <- 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
$ MVector (PrimState (ST (PrimState m))) a
-> Int -> ST (PrimState 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
MVector (PrimState (ST (PrimState m))) a
dataFt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            Int -> a -> m (Int, Int, a)
inner (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (a -> m (Int, Int, a)) -> a -> m (Int, Int, a)
forall a b. (a -> b) -> a -> b
$! a
s a -> a -> a
forall a. Num a => a -> a -> a
- a
ds
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = (Int, Int, a) -> m (Int, Int, a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i, Int
64 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Int
nFt, a
s)
        | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => Int -> a
bit Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
nFt = (Int, Int, a) -> m (Int, Int, a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i, Int
k, a
s)
        | Bool
otherwise = do
            a
t <- 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
$ (a
s +) (a -> a) -> ST (PrimState m) 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) 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 (ST (PrimState m))) a
dataFt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => Int -> a
bit Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            Bool
b <- a -> m Bool
f a
t
            if Bool -> Bool
not Bool
b
              then (Int, Int, a) -> m (Int, Int, a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i, Int
k, a
s)
              else do
                a
di <- 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
$ MVector (PrimState (ST (PrimState m))) a
-> Int -> ST (PrimState 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
MVector (PrimState (ST (PrimState m))) a
dataFt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                Int -> a -> m (Int, Int, a)
inner (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
i)) (a -> m (Int, Int, a)) -> a -> m (Int, Int, a)
forall a b. (a -> b) -> a -> b
$! a
s a -> a -> a
forall a. Num a => a -> a -> a
- a
di
        where
          k :: Int
k = Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

  -- we could start from an arbitrary l, but the API is limited to one
  (!Int
i0, !Int
k0, !a
s0) <- Int -> a -> m (Int, Int, a)
inner Int
l0 (a
0 :: a)

  let inner2 :: Int -> Int -> a -> m Int
inner2 Int
i Int
k_ !a
s
        | Int
k_ Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
        | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => Int -> a
bit Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< MVector (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
VGM.length MVector (PrimState m) a
dataFt = do
            a
t <- 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
$ (a
s +) (a -> a) -> ST (PrimState m) 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) 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 (ST (PrimState m))) a
dataFt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => Int -> a
bit Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            Bool
b <- a -> m Bool
f a
t
            if Bool
b
              then Int -> Int -> a -> m Int
inner2 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => Int -> a
bit Int
k) Int
k a
t
              else Int -> Int -> a -> m Int
inner2 Int
i Int
k a
s
        | Bool
otherwise = Int -> Int -> a -> m Int
inner2 Int
i Int
k a
s
        where
          k :: Int
k = Int
k_ Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

  Int -> Int -> a -> m Int
inner2 Int
i0 Int
k0 a
s0

-- | Applies a binary search on the Fenwick tree. It returns an index \(l\) that satisfies both of
-- the following.
--
-- - \(l = r\) or \(f(a[l] + a[l + 1] + ... + a[r - 1])\) returns `True`.
-- - \(l = 0\) or \(f(a[l - 1] + a[l] + ... + a[r - 1])\) returns `False`.
--
-- If \(f\) is monotone, this is the minimum \(l\) that satisfies
-- \(f(a[l] + a[l + 1] + ... + a[r - 1])\).
--
-- ==== Constraints
--
-- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side
--   effect.
-- - \(f(0)\) returns `True`
-- - \(0 \leq r \leq n\)
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- @since 1.2.2.0
{-# INLINE minLeft #-}
minLeft ::
  (HasCallStack, PrimMonad m, Num a, VU.Unbox a) =>
  -- | The Fenwick tree
  FenwickTree (PrimState m) a ->
  -- | \(r\)
  Int ->
  -- | \(p\): user prediate
  (a -> Bool) ->
  -- | \(l\): \(p\) holds for \([l, r)\)
  m Int
minLeft :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> (a -> Bool) -> m Int
minLeft FenwickTree (PrimState m) a
ft Int
r0 a -> Bool
f = FenwickTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
minLeftM FenwickTree (PrimState m) a
ft Int
r0 (Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> m Bool) -> (a -> Bool) -> a -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
f)

-- | (Extra API) Monadic variant of `minLeft`.
--
-- ==== Constraints
-- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side effect.
-- - \(f(0)\) returns `True`
-- - \(0 \leq l \leq n\)
--
-- ==== Complexity
-- - \(O(\log n)\)
--
-- @since 1.2.2.0
{-# INLINEABLE minLeftM #-}
minLeftM ::
  forall m a.
  (HasCallStack, PrimMonad m, Num a, VU.Unbox a) =>
  -- | The Fenwick tree
  FenwickTree (PrimState m) a ->
  -- | \(r\)
  Int ->
  -- | \(p\): user prediate
  (a -> m Bool) ->
  -- | \(l\): \(p\) holds for \([l, r)\)
  m Int
minLeftM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Num a, Unbox a) =>
FenwickTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
minLeftM FenwickTree {Int
MVector (PrimState m) a
nFt :: forall s a. FenwickTree s a -> Int
dataFt :: forall s a. FenwickTree s a -> MVector s a
nFt :: Int
dataFt :: MVector (PrimState m) a
..} Int
r0 a -> m Bool
f = do
  Bool
b0 <- a -> m Bool
f a
0
  let !()
_ = HasCallStack => Bool -> [Char] -> ()
Bool -> [Char] -> ()
ACIA.runtimeAssert Bool
b0 [Char]
"AtCoder.FenwickTree.minLeftM: `f 0` must return `True`"

  let inner :: Int -> Int -> a -> m (Int, Int, a)
inner Int
i Int
k !a
s
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = (Int, Int, a) -> m (Int, Int, a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i, Int
k, a
s)
        | Bool
otherwise = do
            Bool
b <- a -> m Bool
f a
s
            if Bool -> Bool
not Bool
b
              then (Int, Int, a) -> m (Int, Int, a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i, Int
k, a
s)
              else do
                a
s' <- 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
$ (a
s +) (a -> a) -> ST (PrimState m) 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) 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 (ST (PrimState m))) a
dataFt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                Int -> Int -> a -> m (Int, Int, a)
inner (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
i)) (Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
i) a
s'

  -- we could start from an arbitrary r, but the API is limited to n
  (!Int
i0, !Int
k0, !a
s0) <- Int -> Int -> a -> m (Int, Int, a)
inner Int
r0 (Int
0 :: Int) (a
0 :: a)

  Bool
b0_ <- a -> m Bool
f a
s0
  if Bool
b0_
    then do
      let !()
_ = HasCallStack => Bool -> [Char] -> ()
Bool -> [Char] -> ()
ACIA.runtimeAssert (Int
i0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) [Char]
"AtCoder.FenwickTree.minLeftM: implementation error"
      Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
    else do
      let inner2 :: Int -> Int -> a -> m Int
inner2 Int
i Int
k_ !a
s
            | Int
k_ Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
            | Bool
otherwise = do
                let k :: Int
k = Int
k_ Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                a
t <- 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
$ (a
s -) (a -> a) -> ST (PrimState m) 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) 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 (ST (PrimState m))) a
dataFt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => Int -> a
bit Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                Bool
b <- a -> m Bool
f a
t
                if Bool
b
                  then Int -> Int -> a -> m Int
inner2 Int
i Int
k a
s
                  else Int -> Int -> a -> m Int
inner2 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => Int -> a
bit Int
k) Int
k a
t

      (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> a -> m Int
inner2 Int
i0 Int
k0 a
s0