{-# LANGUAGE RecordWildCards #-}
module AtCoder.SegTree
(
SegTree (nSt, sizeSt, logSt),
new,
build,
write,
modify,
modifyM,
exchange,
read,
prod,
prodMaybe,
allProd,
minLeft,
minLeftM,
maxRight,
maxRightM,
freeze,
unsafeFreeze,
)
where
import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.Internal.Bit qualified as ACIBIT
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Bits (countTrailingZeros, testBit, (.&.), (.>>.))
import Data.Foldable (for_)
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)
data SegTree s a = SegTree
{
forall s a. SegTree s a -> Int
nSt :: {-# UNPACK #-} !Int,
forall s a. SegTree s a -> Int
sizeSt :: {-# UNPACK #-} !Int,
forall s a. SegTree s a -> Int
logSt :: {-# UNPACK #-} !Int,
forall s a. SegTree s a -> MVector s a
dSt :: !(VUM.MVector s a)
}
{-# INLINE new #-}
new :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Int -> m (SegTree (PrimState m) a)
new :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Int -> m (SegTree (PrimState m) a)
new Int
nSt
| Int
nSt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 = Vector a -> m (SegTree (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vector a -> m (SegTree (PrimState m) a)
build (Vector a -> m (SegTree (PrimState m) a))
-> Vector a -> m (SegTree (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Int -> a -> Vector a
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nSt a
forall a. Monoid a => a
mempty
| Bool
otherwise = [Char] -> m (SegTree (PrimState m) a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> m (SegTree (PrimState m) a))
-> [Char] -> m (SegTree (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ [Char]
"AtCoder.SegTree.new: given negative size (`" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
nSt [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"`)"
{-# INLINE build #-}
build :: (PrimMonad m, Monoid a, VU.Unbox a) => VU.Vector a -> m (SegTree (PrimState m) a)
build :: forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vector a -> m (SegTree (PrimState m) a)
build Vector a
vs = do
let nSt :: Int
nSt = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
vs
let sizeSt :: Int
sizeSt = Int -> Int
ACIBIT.bitCeil Int
nSt
let logSt :: Int
logSt = Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
sizeSt
MVector (PrimState m) a
dSt <- Int -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
sizeSt) a
forall a. Monoid a => a
mempty
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
vs ((Int -> a -> m ()) -> m ()) -> (Int -> a -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i a
v -> do
MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
dSt (Int
sizeSt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) a
v
let segtree :: SegTree (PrimState m) a
segtree = SegTree {Int
MVector (PrimState m) a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..}
[Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
sizeSt Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
sizeSt Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
1] ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
SegTree (PrimState m) a -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> m ()
update SegTree (PrimState m) a
segtree Int
i
SegTree (PrimState m) a -> m (SegTree (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SegTree (PrimState m) a
segtree
{-# INLINE write #-}
write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> a -> m ()
write :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m ()
write self :: SegTree (PrimState m) a
self@SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} Int
p a
x = do
let !()
_ = HasCallStack => [Char] -> Int -> Int -> ()
[Char] -> Int -> Int -> ()
ACIA.checkIndex [Char]
"AtCoder.SegTree.write" Int
p Int
nSt
MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
dSt (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) a
x
[Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
1 .. Int
logSt] ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
SegTree (PrimState m) a -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> m ()
update SegTree (PrimState m) a
self ((Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
i)
{-# INLINE modify #-}
modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> (a -> a) -> Int -> m ()
modify :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> a) -> Int -> m ()
modify self :: SegTree (PrimState m) a
self@SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} a -> a
f Int
p = do
let !()
_ = HasCallStack => [Char] -> Int -> Int -> ()
[Char] -> Int -> Int -> ()
ACIA.checkIndex [Char]
"AtCoder.SegTree.modify" Int
p Int
nSt
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
dSt a -> a
f (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt)
[Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
1 .. Int
logSt] ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
SegTree (PrimState m) a -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> m ()
update SegTree (PrimState m) a
self ((Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
i)
{-# INLINE modifyM #-}
modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM self :: SegTree (PrimState m) a
self@SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} a -> m a
f Int
p = do
let !()
_ = HasCallStack => [Char] -> Int -> Int -> ()
[Char] -> Int -> Int -> ()
ACIA.checkIndex [Char]
"AtCoder.SegTree.modifyM" Int
p Int
nSt
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
dSt a -> m a
f (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt)
[Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
1 .. Int
logSt] ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
SegTree (PrimState m) a -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> m ()
update SegTree (PrimState m) a
self ((Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
i)
{-# INLINE exchange #-}
exchange :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> a -> m a
exchange :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m a
exchange self :: SegTree (PrimState m) a
self@SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} Int
p a
x = do
let !()
_ = HasCallStack => [Char] -> Int -> Int -> ()
[Char] -> Int -> Int -> ()
ACIA.checkIndex [Char]
"AtCoder.SegTree.exchange" Int
p Int
nSt
a
ret <- MVector (PrimState m) a -> Int -> a -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector (PrimState m) a
dSt (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) a
x
MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
dSt (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) a
x
[Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
1 .. Int
logSt] ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
SegTree (PrimState m) a -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> m ()
update SegTree (PrimState m) a
self ((Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
i)
a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
ret
{-# INLINE read #-}
read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> m a
read :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> m a
read SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} Int
p = do
let !()
_ = HasCallStack => [Char] -> Int -> Int -> ()
[Char] -> Int -> Int -> ()
ACIA.checkIndex [Char]
"AtCoder.SegTree.read" Int
p Int
nSt
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
dSt (Int -> m a) -> Int -> m a
forall a b. (a -> b) -> a -> b
$ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt
{-# INLINE prod #-}
prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> Int -> m a
prod :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
prod self :: SegTree (PrimState m) a
self@SegTree {Int
nSt :: forall s a. SegTree s a -> Int
nSt :: Int
nSt} Int
l0 Int
r0
| Int -> Int -> Int -> Bool
ACIA.testInterval Int
l0 Int
r0 Int
nSt = SegTree (PrimState m) a -> Int -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
unsafeProd SegTree (PrimState m) a
self Int
l0 Int
r0
| Bool
otherwise = [Char] -> Int -> Int -> Int -> m a
forall a. HasCallStack => [Char] -> Int -> Int -> Int -> a
ACIA.errorInterval [Char]
"AtCoder.SegTree.prod" Int
l0 Int
r0 Int
nSt
{-# INLINE prodMaybe #-}
prodMaybe :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> Int -> m (Maybe a)
prodMaybe :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m (Maybe a)
prodMaybe self :: SegTree (PrimState m) a
self@SegTree {Int
nSt :: forall s a. SegTree s a -> Int
nSt :: Int
nSt} Int
l0 Int
r0
| Int -> Int -> Int -> Bool
ACIA.testInterval Int
l0 Int
r0 Int
nSt = 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
<$> SegTree (PrimState m) a -> Int -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
unsafeProd SegTree (PrimState m) a
self Int
l0 Int
r0
| 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
{-# INLINE unsafeProd #-}
unsafeProd :: (PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> Int -> m a
unsafeProd :: forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
unsafeProd SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} Int
l0 Int
r0 = Int -> Int -> a -> a -> m a
inner (Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) (Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a
forall a. Monoid a => a
mempty a
forall a. Monoid a => a
mempty
where
inner :: Int -> Int -> a -> a -> m a
inner Int
l Int
r !a
smL !a
smR
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 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
smL a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
smR
| Bool
otherwise = do
!a
smL' <-
if Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
l Int
0
then (a
smL <>) (a -> a) -> m a -> m 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
dSt Int
l
else a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
smL
!a
smR' <-
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
r Int
0
then (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
smR) (a -> a) -> m a -> m 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
dSt Int
r
else a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
smR
Int -> Int -> a -> a -> m a
inner ((Int
l 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
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
1) a
smL' a
smR'
{-# INLINE allProd #-}
allProd :: (PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> m a
allProd :: forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> m a
allProd SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} = 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
dSt Int
1
{-# INLINE minLeft #-}
minLeft ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
SegTree (PrimState m) a ->
Int ->
(a -> Bool) ->
m Int
minLeft :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> (a -> Bool) -> m Int
minLeft SegTree (PrimState m) a
seg Int
r0 a -> Bool
f = SegTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
minLeftM SegTree (PrimState m) a
seg 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)
{-# INLINE minLeftM #-}
minLeftM ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
SegTree (PrimState m) a ->
Int ->
(a -> m Bool) ->
m Int
minLeftM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
minLeftM SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} Int
r0 a -> m Bool
f = do
Bool
b <- a -> m Bool
f a
forall a. Monoid a => a
mempty
let !()
_ = HasCallStack => Bool -> [Char] -> ()
Bool -> [Char] -> ()
ACIA.runtimeAssert Bool
b [Char]
"AtCoder.SegTree.minLeftM: `f empty` returned `False`"
if Int
r0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
then Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
else Int -> a -> m Int
inner (Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) a
forall a. Monoid a => a
mempty
where
!()
_ = HasCallStack => Bool -> [Char] -> ()
Bool -> [Char] -> ()
ACIA.runtimeAssert (Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r0 Bool -> Bool -> Bool
&& Int
r0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
nSt) ([Char] -> ()) -> [Char] -> ()
forall a b. (a -> b) -> a -> b
$ [Char]
"AtCoder.SegTree.minLeftM: given invalid `right` index `" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
r0 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"` over length `" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
nSt [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"`"
inner :: Int -> a -> m Int
inner Int
r !a
sm = do
let r' :: Int
r' = Int -> Int
forall {b}. (Integral b, Bits b) => b -> b
chooseBit (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
!a
sm' <- (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
sm) (a -> a) -> m a -> m 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
dSt Int
r'
Bool
b <- a -> m Bool
f a
sm'
if Bool -> Bool
not Bool
b
then do
Int -> a -> m Int
inner2 Int
r' a
sm
else do
if (Int
r' Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
r')) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
r'
then Int -> a -> m Int
inner Int
r' a
sm'
else Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
chooseBit :: b -> b
chooseBit b
r
| b
r b -> b -> Bool
forall a. Ord a => a -> a -> Bool
> b
1 Bool -> Bool -> Bool
&& b -> Bool
forall a. Integral a => a -> Bool
odd b
r = b -> b
chooseBit (b -> b) -> b -> b
forall a b. (a -> b) -> a -> b
$ b
r b -> Int -> b
forall a. Bits a => a -> Int -> a
.>>. Int
1
| Bool
otherwise = b
r
inner2 :: Int -> a -> m Int
inner2 Int
r a
sm
| Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sizeSt = do
let r' :: Int
r' = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
!a
sm' <- (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
sm) (a -> a) -> m a -> m 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
dSt Int
r'
Bool
b <- a -> m Bool
f a
sm'
if Bool
b
then Int -> a -> m Int
inner2 (Int
r' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a
sm'
else Int -> a -> m Int
inner2 Int
r' a
sm
| Bool
otherwise = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
sizeSt
{-# INLINE maxRight #-}
maxRight ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
SegTree (PrimState m) a ->
Int ->
(a -> Bool) ->
m Int
maxRight :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> (a -> Bool) -> m Int
maxRight SegTree (PrimState m) a
seg Int
l0 a -> Bool
f = SegTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
maxRightM SegTree (PrimState m) a
seg 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)
{-# INLINE maxRightM #-}
maxRightM ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
SegTree (PrimState m) a ->
Int ->
(a -> m Bool) ->
m Int
maxRightM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> (a -> m Bool) -> m Int
maxRightM SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} Int
l0 a -> m Bool
f = do
Bool
b <- a -> m Bool
f a
forall a. Monoid a => a
mempty
let !()
_ = HasCallStack => Bool -> [Char] -> ()
Bool -> [Char] -> ()
ACIA.runtimeAssert Bool
b [Char]
"AtCoder.SegTree.maxRightM: `f mempty` returned `False`"
if Int
l0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nSt
then Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
nSt
else Int -> a -> m Int
inner (Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt) a
forall a. Monoid a => a
mempty
where
!()
_ = HasCallStack => Bool -> [Char] -> ()
Bool -> [Char] -> ()
ACIA.runtimeAssert (Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l0 Bool -> Bool -> Bool
&& Int
l0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
nSt) ([Char] -> ()) -> [Char] -> ()
forall a b. (a -> b) -> a -> b
$ [Char]
"AtCoder.SegTree.maxRightM: given invalid `left` index `" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
l0 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"` over length `" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
nSt [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"`"
inner :: Int -> a -> m Int
inner Int
l !a
sm = do
let l' :: Int
l' = Int -> Int
chooseBit Int
l
!a
sm' <- (a
sm <>) (a -> a) -> m a -> m 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
dSt Int
l'
Bool
b <- a -> m Bool
f a
sm'
if Bool -> Bool
not Bool
b
then do
Int -> a -> m Int
inner2 Int
l' a
sm
else do
let l'' :: Int
l'' = Int
l' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
if (Int
l'' Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
l'')) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
l''
then Int -> a -> m Int
inner Int
l'' a
sm'
else Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
nSt
chooseBit :: Int -> Int
chooseBit :: Int -> Int
chooseBit Int
l
| Int -> Bool
forall a. Integral a => a -> Bool
even Int
l = Int -> Int
chooseBit (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int
l Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
1
| Bool
otherwise = Int
l
inner2 :: Int -> a -> m Int
inner2 Int
l !a
sm
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sizeSt = do
let l' :: Int
l' = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
l
!a
sm' <- (a
sm <>) (a -> a) -> m a -> m 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
dSt Int
l'
Bool
b <- a -> m Bool
f a
sm'
if Bool
b
then Int -> a -> m Int
inner2 (Int
l' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
sm'
else Int -> a -> m Int
inner2 Int
l' a
sm
| Bool
otherwise = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
sizeSt
{-# INLINE freeze #-}
freeze :: (PrimMonad m, VU.Unbox a) => SegTree (PrimState m) a -> m (VU.Vector a)
freeze :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
SegTree (PrimState m) a -> m (Vector a)
freeze SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} = do
MVector (PrimState m) a -> m (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.freeze (MVector (PrimState m) a -> m (Vector a))
-> (MVector (PrimState m) a -> MVector (PrimState m) a)
-> MVector (PrimState m) a
-> m (Vector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.take Int
nSt (MVector (PrimState m) a -> m (Vector a))
-> MVector (PrimState m) a -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.drop Int
sizeSt MVector (PrimState m) a
dSt
{-# INLINE unsafeFreeze #-}
unsafeFreeze :: (PrimMonad m, VU.Unbox a) => SegTree (PrimState m) a -> m (VU.Vector a)
unsafeFreeze :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
SegTree (PrimState m) a -> m (Vector a)
unsafeFreeze SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} = do
MVector (PrimState m) a -> m (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (MVector (PrimState m) a -> m (Vector a))
-> (MVector (PrimState m) a -> MVector (PrimState m) a)
-> MVector (PrimState m) a
-> m (Vector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.take Int
nSt (MVector (PrimState m) a -> m (Vector a))
-> MVector (PrimState m) a -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.drop Int
sizeSt MVector (PrimState m) a
dSt
{-# INLINE update #-}
update :: (PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> m ()
update :: forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> m ()
update SegTree {Int
MVector (PrimState m) a
nSt :: forall s a. SegTree s a -> Int
sizeSt :: forall s a. SegTree s a -> Int
logSt :: forall s a. SegTree s a -> Int
dSt :: forall s a. SegTree s a -> MVector s a
nSt :: Int
sizeSt :: Int
logSt :: Int
dSt :: MVector (PrimState m) a
..} Int
k = do
a
opL <- 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
dSt (Int -> m a) -> Int -> m a
forall a b. (a -> b) -> a -> b
$ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k
a
opR <- 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
dSt (Int -> m a) -> Int -> m a
forall a b. (a -> b) -> a -> b
$ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
dSt Int
k (a -> m ()) -> a -> m ()
forall a b. (a -> b) -> a -> b
$! a
opL a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
opR