{-# LANGUAGE RecordWildCards #-}
module AtCoder.Extra.SegTree2d
(
SegTree2d (..),
new,
build,
build2,
build3,
write,
modify,
modifyM,
prod,
allProd,
count,
)
where
import AtCoder.Extra.Bisect (lowerBound)
import AtCoder.Extra.Vector (argsort)
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.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)
data SegTree2d s a = SegTree2d
{
forall s a. SegTree2d s a -> Int
nSt :: {-# UNPACK #-} !Int,
forall s a. SegTree2d s a -> Int
nxSt :: {-# UNPACK #-} !Int,
forall s a. SegTree2d s a -> Int
logSt :: {-# UNPACK #-} !Int,
forall s a. SegTree2d s a -> Int
sizeSt :: {-# UNPACK #-} !Int,
forall s a. SegTree2d s a -> Vector Int
dictXSt :: {-# UNPACK #-} !(VU.Vector Int),
forall s a. SegTree2d s a -> Vector Int
allYSt :: !(VU.Vector Int),
forall s a. SegTree2d s a -> Vector Int
posSt :: !(VU.Vector Int),
forall s a. SegTree2d s a -> Vector Int
indptrSt :: !(VU.Vector Int),
forall s a. SegTree2d s a -> MVector s a
dataSt :: !(VUM.MVector s a),
forall s a. SegTree2d s a -> Vector Int
toLeftSt :: !(VU.Vector Int)
}
{-# INLINEABLE new #-}
new ::
(PrimMonad m, Monoid a, VU.Unbox a) =>
VU.Vector (Int, Int) ->
m (SegTree2d (PrimState m) a)
new :: forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vector (Int, Int) -> m (SegTree2d (PrimState m) a)
new Vector (Int, Int)
xys = ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a))
-> ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Vector Int
-> Vector Int
-> Vector a
-> ST (PrimState m) (SegTree2d (PrimState m) a)
forall s a.
(HasCallStack, Monoid a, Unbox a) =>
Vector Int -> Vector Int -> Vector a -> ST s (SegTree2d s a)
buildST Vector Int
xs Vector Int
ys (Int -> a -> Vector a
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
n a
forall a. Monoid a => a
mempty)
where
n :: Int
n = Vector (Int, Int) -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector (Int, Int)
xys
(!Vector Int
xs, !Vector Int
ys) = Vector (Int, Int) -> (Vector Int, Vector Int)
forall a b.
(Unbox a, Unbox b) =>
Vector (a, b) -> (Vector a, Vector b)
VU.unzip Vector (Int, Int)
xys
{-# INLINE build #-}
build ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
VU.Vector Int ->
VU.Vector Int ->
VU.Vector a ->
m (SegTree2d (PrimState m) a)
build :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Vector Int
-> Vector Int -> Vector a -> m (SegTree2d (PrimState m) a)
build Vector Int
xs Vector Int
ys Vector a
ws = ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a))
-> ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Vector Int
-> Vector Int
-> Vector a
-> ST (PrimState m) (SegTree2d (PrimState m) a)
forall s a.
(HasCallStack, Monoid a, Unbox a) =>
Vector Int -> Vector Int -> Vector a -> ST s (SegTree2d s a)
buildST Vector Int
xs Vector Int
ys Vector a
ws
{-# INLINE build2 #-}
build2 ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
VU.Vector (Int, Int) ->
VU.Vector a ->
m (SegTree2d (PrimState m) a)
build2 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Vector (Int, Int) -> Vector a -> m (SegTree2d (PrimState m) a)
build2 Vector (Int, Int)
xys Vector a
ws = ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a))
-> ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ do
let (!Vector Int
xs, !Vector Int
ys) = Vector (Int, Int) -> (Vector Int, Vector Int)
forall a b.
(Unbox a, Unbox b) =>
Vector (a, b) -> (Vector a, Vector b)
VU.unzip Vector (Int, Int)
xys
Vector Int
-> Vector Int
-> Vector a
-> ST (PrimState m) (SegTree2d (PrimState m) a)
forall s a.
(HasCallStack, Monoid a, Unbox a) =>
Vector Int -> Vector Int -> Vector a -> ST s (SegTree2d s a)
buildST Vector Int
xs Vector Int
ys Vector a
ws
{-# INLINE build3 #-}
build3 :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => VU.Vector (Int, Int, a) -> m (SegTree2d (PrimState m) a)
build3 :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Vector (Int, Int, a) -> m (SegTree2d (PrimState m) a)
build3 Vector (Int, Int, a)
xyws = ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a))
-> ST (PrimState m) (SegTree2d (PrimState m) a)
-> m (SegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ do
let (!Vector Int
xs, !Vector Int
ys, !Vector a
ws) = Vector (Int, Int, a) -> (Vector Int, Vector Int, Vector a)
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
Vector (a, b, c) -> (Vector a, Vector b, Vector c)
VU.unzip3 Vector (Int, Int, a)
xyws
Vector Int
-> Vector Int
-> Vector a
-> ST (PrimState m) (SegTree2d (PrimState m) a)
forall s a.
(HasCallStack, Monoid a, Unbox a) =>
Vector Int -> Vector Int -> Vector a -> ST s (SegTree2d s a)
buildST Vector Int
xs Vector Int
ys Vector a
ws
{-# INLINE write #-}
write ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
SegTree2d (PrimState m) a ->
Int ->
a ->
m ()
write :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> Int -> a -> m ()
write SegTree2d (PrimState m) a
seg Int
i a
x = 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
SegTree2d (PrimState (ST (PrimState m))) a
-> (a -> ST (PrimState m) a) -> Int -> ST (PrimState m) ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM SegTree2d (PrimState m) a
SegTree2d (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
x) Int
i
{-# INLINE modify #-}
modify ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
SegTree2d (PrimState m) a ->
(a -> a) ->
Int ->
m ()
modify :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> (a -> a) -> Int -> m ()
modify SegTree2d (PrimState m) a
seg a -> a
f Int
i = 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
SegTree2d (PrimState (ST (PrimState m))) a
-> (a -> ST (PrimState m) a) -> Int -> ST (PrimState m) ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM SegTree2d (PrimState m) a
SegTree2d (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
i
{-# INLINEABLE modifyM #-}
modifyM ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
SegTree2d (PrimState m) a ->
(a -> m a) ->
Int ->
m ()
modifyM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM seg :: SegTree2d (PrimState m) a
seg@SegTree2d {Int
MVector (PrimState m) a
Vector Int
nSt :: forall s a. SegTree2d s a -> Int
nxSt :: forall s a. SegTree2d s a -> Int
logSt :: forall s a. SegTree2d s a -> Int
sizeSt :: forall s a. SegTree2d s a -> Int
dictXSt :: forall s a. SegTree2d s a -> Vector Int
allYSt :: forall s a. SegTree2d s a -> Vector Int
posSt :: forall s a. SegTree2d s a -> Vector Int
indptrSt :: forall s a. SegTree2d s a -> Vector Int
dataSt :: forall s a. SegTree2d s a -> MVector s a
toLeftSt :: forall s a. SegTree2d s a -> Vector Int
nSt :: Int
nxSt :: Int
logSt :: Int
sizeSt :: Int
dictXSt :: Vector Int
allYSt :: Vector Int
posSt :: Vector Int
indptrSt :: Vector Int
dataSt :: MVector (PrimState m) a
toLeftSt :: Vector Int
..} a -> m a
f Int
rawIdx = do
let !()
_ = HasCallStack => String -> Int -> Int -> ()
String -> Int -> Int -> ()
ACIA.checkIndex String
"AtCoder.Extra.SegTree2d.modifyM" Int
rawIdx Int
nSt
Int -> Int -> m ()
inner Int
1 (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Vector Int
posSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
rawIdx
where
inner :: Int -> Int -> m ()
inner Int
i Int
p = do
let indptrI :: Int
indptrI = Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
SegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyIST SegTree2d (PrimState m) a
seg a -> m a
f Int
i (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
indptrI
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sizeSt) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
let lc :: Int
lc = Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
indptrI
let rc :: Int
rc = (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
indptrI) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lc
if Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
then do
let i' :: Int
i' = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0
let p' :: Int
p' = Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lc
Int -> Int -> m ()
inner Int
i' Int
p'
else do
let i' :: Int
i' = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
let p' :: Int
p' = Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rc
Int -> Int -> m ()
inner Int
i' Int
p'
{-# INLINE prod #-}
prod ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
SegTree2d (PrimState m) a ->
Int ->
Int ->
Int ->
Int ->
m a
prod :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a
prod SegTree2d (PrimState m) a
seg Int
lx Int
rx Int
ly Int
ry = 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
$ SegTree2d (PrimState m) a
-> Int -> Int -> Int -> Int -> ST (PrimState m) a
forall s a.
(HasCallStack, Monoid a, Unbox a) =>
SegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
prodST SegTree2d (PrimState m) a
seg Int
lx Int
rx Int
ly Int
ry
{-# INLINE allProd #-}
allProd :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree2d (PrimState m) a -> m a
allProd :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> m a
allProd SegTree2d (PrimState m) a
seg = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
forall a. Monoid a => a
mempty (Maybe a -> a) -> m (Maybe a) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) a -> Int -> m (Maybe a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (Maybe a)
VGM.readMaybe (SegTree2d (PrimState m) a -> MVector (PrimState m) a
forall s a. SegTree2d s a -> MVector s a
dataSt SegTree2d (PrimState m) a
seg) Int
1
{-# INLINE count #-}
count ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
SegTree2d (PrimState m) a ->
Int ->
Int ->
Int ->
Int ->
m Int
count :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> Int -> Int -> Int -> Int -> m Int
count SegTree2d (PrimState m) a
seg Int
lx Int
rx Int
ly Int
ry = ST (PrimState m) Int -> m Int
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Int -> m Int) -> ST (PrimState m) Int -> m Int
forall a b. (a -> b) -> a -> b
$ SegTree2d (PrimState m) a
-> Int -> Int -> Int -> Int -> ST (PrimState m) Int
forall s a.
(HasCallStack, Monoid a, Unbox a) =>
SegTree2d s a -> Int -> Int -> Int -> Int -> ST s Int
countST SegTree2d (PrimState m) a
seg Int
lx Int
rx Int
ly Int
ry
{-# INLINEABLE buildST #-}
buildST :: forall s a. (HasCallStack, Monoid a, VU.Unbox a) => VU.Vector Int -> VU.Vector Int -> VU.Vector a -> ST s (SegTree2d s a)
buildST :: forall s a.
(HasCallStack, Monoid a, Unbox a) =>
Vector Int -> Vector Int -> Vector a -> ST s (SegTree2d s a)
buildST Vector Int
xs Vector Int
ys Vector a
ws = do
let nSt :: Int
nSt = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
xs
let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
nSt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
ys Bool -> Bool -> Bool
&& Int
nSt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
ws) String
"AtCoder.Extra.SegTree2d.buildST: length mismatch among `xs`, `ys` and `ws`"
let dictXSt :: Vector Int
dictXSt = Vector Int -> Vector Int
forall a. (Unbox a, Eq a) => Vector a -> Vector a
VU.uniq (Vector Int -> Vector Int) -> Vector Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ (forall s. MVector s Int -> ST s ()) -> Vector Int -> Vector Int
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify MVector s Int -> ST s ()
MVector (PrimState (ST s)) Int -> ST s ()
forall s. MVector s Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> m ()
VAI.sort Vector Int
xs
let nxSt :: Int
nxSt = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
dictXSt
let logSt :: Int
logSt = Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Int
ACIB.bitCeil (Int
nxSt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
let sizeSt :: Int
sizeSt = Int -> Int
forall a. Bits a => Int -> a
bit Int
logSt
let compressedXs :: Vector Int
compressedXs = (Int -> Int) -> Vector Int -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> (Int -> Maybe Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int -> Maybe Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Maybe Int
lowerBound Vector Int
dictXSt) Vector Int
xs
let indptrSt :: Vector Int
indptrSt = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
MVector (PrimState (ST s)) Int
indptr <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
0 :: Int)
Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector Int
compressedXs ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
let inner :: Int -> ST s ()
inner Int
j
| Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 = do
MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState (ST s)) Int
indptr (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
Int -> ST s ()
inner (Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
| Bool
otherwise = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
Int -> ST s ()
inner (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt
MVector (PrimState (ST s)) Int
-> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (Int -> a -> m b) -> m ()
VUM.iforM_ (MVector s Int -> MVector s Int
forall a s. Unbox a => MVector s a -> MVector s a
VUM.init MVector s Int
MVector (PrimState (ST s)) Int
indptr) ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i Int
x -> do
MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState (ST s)) Int
indptr (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
MVector (PrimState (ST s)) Int
indptr
MVector (PrimState (ST s)) a
dataSt <- 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
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Vector Int -> Int
forall a. Unbox a => Vector a -> a
VU.last Vector Int
indptrSt) (a
forall a. Monoid a => a
mempty :: a)
let yis :: Vector Int
yis = Vector Int -> Vector Int
forall a. (Ord a, Unbox a) => Vector a -> Vector Int
argsort Vector Int
ys
let posSt :: Vector Int
posSt = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
MVector s Int
vec <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nSt (Int
0 :: Int)
Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
yis ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i Int
yi -> do
MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
vec Int
yi Int
i
MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
vec
Vector Int
toLeftSt <- do
MVector (PrimState (ST s)) Int
toLeft <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
sizeSt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
0 :: Int)
MVector (PrimState (ST s)) Int
ptr <- Vector Int -> ST s (MVector (PrimState (ST s)) Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw Vector Int
indptrSt
Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector Int
yis ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
rawIdx -> do
let inner :: Int -> Int -> ST s ()
inner Int
i Int
j
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = do
Int
p <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState (ST s)) Int
ptr Int
i
MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
ptr Int
i (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
p 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 (PrimState (ST s)) a
dataSt (Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p) (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$ Vector a
ws Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
rawIdx
Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= -Int
1) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
toLeft (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ if Int -> Bool
forall a. Integral a => a -> Bool
even Int
j then Int
1 else Int
0
Int -> Int -> ST s ()
inner (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int
i
| Bool
otherwise = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
let i0 :: Int
i0 = Vector Int
compressedXs Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
rawIdx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeSt
Int -> Int -> ST s ()
inner Int
i0 (-Int
1)
MVector (PrimState (ST s)) Int
-> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (Int -> a -> m b) -> m ()
VUM.iforM_ (MVector s Int -> MVector s Int
forall a s. Unbox a => MVector s a -> MVector s a
VUM.init MVector s Int
MVector (PrimState (ST s)) Int
toLeft) ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i Int
x -> do
MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState (ST s)) Int
toLeft (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState (ST s)) Int
toLeft
[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
sizeSt 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
i -> do
let off :: Int
off = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
let n :: Int
n = Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
[Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> 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 (PrimState (ST s)) a
dataSt (Int -> ST s a) -> Int -> ST s a
forall a b. (a -> b) -> a -> b
$ Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j 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 (PrimState (ST s)) a
dataSt (Int -> ST s a) -> Int -> ST s a
forall a b. (a -> b) -> a -> b
$ Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j 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 (PrimState (ST s)) a
dataSt (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) (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
let allYSt :: Vector Int
allYSt = (forall s. MVector s Int -> ST s ()) -> Vector Int -> Vector Int
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify MVector s Int -> ST s ()
MVector (PrimState (ST s)) Int -> ST s ()
forall s. MVector s Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> m ()
VAI.sort Vector Int
ys
SegTree2d s a -> ST s (SegTree2d s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SegTree2d {Int
MVector s a
MVector (PrimState (ST s)) a
Vector Int
nSt :: Int
nxSt :: Int
logSt :: Int
sizeSt :: Int
dictXSt :: Vector Int
allYSt :: Vector Int
posSt :: Vector Int
indptrSt :: Vector Int
dataSt :: MVector s a
toLeftSt :: Vector Int
nSt :: Int
dictXSt :: Vector Int
nxSt :: Int
logSt :: Int
sizeSt :: Int
indptrSt :: Vector Int
dataSt :: MVector (PrimState (ST s)) a
posSt :: Vector Int
toLeftSt :: Vector Int
allYSt :: Vector Int
..}
{-# INLINEABLE modifyIST #-}
modifyIST :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyIST :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyIST SegTree2d {Int
MVector (PrimState m) a
Vector Int
nSt :: forall s a. SegTree2d s a -> Int
nxSt :: forall s a. SegTree2d s a -> Int
logSt :: forall s a. SegTree2d s a -> Int
sizeSt :: forall s a. SegTree2d s a -> Int
dictXSt :: forall s a. SegTree2d s a -> Vector Int
allYSt :: forall s a. SegTree2d s a -> Vector Int
posSt :: forall s a. SegTree2d s a -> Vector Int
indptrSt :: forall s a. SegTree2d s a -> Vector Int
dataSt :: forall s a. SegTree2d s a -> MVector s a
toLeftSt :: forall s a. SegTree2d s a -> Vector Int
nSt :: Int
nxSt :: Int
logSt :: Int
sizeSt :: Int
dictXSt :: Vector Int
allYSt :: Vector Int
posSt :: Vector Int
indptrSt :: Vector Int
dataSt :: MVector (PrimState m) a
toLeftSt :: Vector Int
..} a -> m a
f Int
i Int
j0 = do
let lid :: Int
lid = Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
let off :: Int
off = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
lid
let inner :: Int -> f ()
inner Int
j_ = do
Bool -> f () -> f ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j_ Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) (f () -> f ()) -> f () -> f ()
forall a b. (a -> b) -> a -> b
$ do
let j :: Int
j = Int
j_ Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
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
dataSt (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j 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
dataSt (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j 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
dataSt (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) (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 ()
inner Int
j
let j :: Int
j = Int
j0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lid
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
dataSt a -> m a
f (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)
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
$ Int -> ST (PrimState m) ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> f ()
inner Int
j
{-# INLINEABLE prodST #-}
prodST :: forall s a. (HasCallStack, Monoid a, VU.Unbox a) => SegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
prodST :: forall s a.
(HasCallStack, Monoid a, Unbox a) =>
SegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
prodST seg :: SegTree2d s a
seg@SegTree2d {Int
MVector s a
Vector Int
nSt :: forall s a. SegTree2d s a -> Int
nxSt :: forall s a. SegTree2d s a -> Int
logSt :: forall s a. SegTree2d s a -> Int
sizeSt :: forall s a. SegTree2d s a -> Int
dictXSt :: forall s a. SegTree2d s a -> Vector Int
allYSt :: forall s a. SegTree2d s a -> Vector Int
posSt :: forall s a. SegTree2d s a -> Vector Int
indptrSt :: forall s a. SegTree2d s a -> Vector Int
dataSt :: forall s a. SegTree2d s a -> MVector s a
toLeftSt :: forall s a. SegTree2d s a -> Vector Int
nSt :: Int
nxSt :: Int
logSt :: Int
sizeSt :: Int
dictXSt :: Vector Int
allYSt :: Vector Int
posSt :: Vector Int
indptrSt :: Vector Int
dataSt :: MVector s a
toLeftSt :: Vector Int
..} Int
lx Int
rx Int
ly Int
ry = do
let a0 :: Int
a0 = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
allYSt) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> Maybe Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Maybe Int
lowerBound Vector Int
allYSt Int
ly
let b0 :: Int
b0 = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
allYSt) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> Maybe Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Maybe Int
lowerBound Vector Int
allYSt Int
ry
a -> Int -> Int -> Int -> Int -> Int -> ST s a
dfs a
forall a. Monoid a => a
mempty Int
1 Int
0 Int
sizeSt Int
a0 Int
b0
where
!()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
lx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
rx Bool -> Bool -> Bool
&& Int
ly Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
ry) String
"AtCoder.Extra.SegTree2d.prodST: given invalid rectangle"
!l0 :: Int
l0 = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
dictXSt) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> Maybe Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Maybe Int
lowerBound Vector Int
dictXSt Int
lx
!r0 :: Int
r0 = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
dictXSt) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> Maybe Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Maybe Int
lowerBound Vector Int
dictXSt Int
rx
dfs :: a -> Int -> Int -> Int -> Int -> Int -> ST s a
dfs :: a -> Int -> Int -> Int -> Int -> Int -> ST s a
dfs !a
res Int
i Int
l Int
r Int
a Int
b
| Int
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
b = a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res
| Int
r0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l Bool -> Bool -> Bool
|| Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l0 = a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res
| Int
l0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l Bool -> Bool -> Bool
&& Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r0 = do
a
xi <- SegTree2d s a -> Int -> Int -> Int -> ST s a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
SegTree2d s a -> Int -> Int -> Int -> ST s a
prodIST SegTree2d s a
seg Int
i Int
a Int
b
a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST s a) -> a -> ST s a
forall a b. (a -> b) -> a -> b
$! a
xi a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
res
| Bool
otherwise = do
let indptrI :: Int
indptrI = Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
let toLeftI :: Int
toLeftI = Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
indptrI
let la :: Int
la = Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
indptrI Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
a) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
toLeftI
let ra :: Int
ra = Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
la
let lb :: Int
lb = Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
indptrI Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
toLeftI
let rb :: Int
rb = Int
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lb
let m :: Int
m = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
a
res' <- a -> Int -> Int -> Int -> Int -> Int -> ST s a
dfs a
res (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Int
l Int
m Int
la Int
lb
a -> Int -> Int -> Int -> Int -> Int -> ST s a
dfs a
res' (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
m Int
r Int
ra Int
rb
{-# INLINEABLE prodIST #-}
prodIST :: (HasCallStack, Monoid a, VU.Unbox a) => SegTree2d s a -> Int -> Int -> Int -> ST s a
prodIST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
SegTree2d s a -> Int -> Int -> Int -> ST s a
prodIST SegTree2d {Int
MVector s a
Vector Int
nSt :: forall s a. SegTree2d s a -> Int
nxSt :: forall s a. SegTree2d s a -> Int
logSt :: forall s a. SegTree2d s a -> Int
sizeSt :: forall s a. SegTree2d s a -> Int
dictXSt :: forall s a. SegTree2d s a -> Vector Int
allYSt :: forall s a. SegTree2d s a -> Vector Int
posSt :: forall s a. SegTree2d s a -> Vector Int
indptrSt :: forall s a. SegTree2d s a -> Vector Int
dataSt :: forall s a. SegTree2d s a -> MVector s a
toLeftSt :: forall s a. SegTree2d s a -> Vector Int
nSt :: Int
nxSt :: Int
logSt :: Int
sizeSt :: Int
dictXSt :: Vector Int
allYSt :: Vector Int
posSt :: Vector Int
indptrSt :: Vector Int
dataSt :: MVector s a
toLeftSt :: Vector Int
..} Int
i Int
a Int
b = a -> Int -> Int -> ST s a
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
a -> Int -> Int -> m a
inner a
forall a. Monoid a => a
mempty (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
a) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
where
lid :: Int
lid = Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
off :: Int
off = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
lid
n :: Int
n = Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lid
inner :: a -> Int -> Int -> m a
inner !a
res Int
l Int
r
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r = do
a
res' <-
if Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
l Int
0
then (a
res <>) (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 s a
MVector (PrimState m) a
dataSt (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l)
else a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res
a
res'' <-
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
res') (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 s a
MVector (PrimState m) a
dataSt (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r)
else a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res'
a -> Int -> Int -> m a
inner a
res'' ((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)
| Bool
otherwise = a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res
{-# INLINEABLE countST #-}
countST :: forall s a. (HasCallStack, Monoid a, VU.Unbox a) => SegTree2d s a -> Int -> Int -> Int -> Int -> ST s Int
countST :: forall s a.
(HasCallStack, Monoid a, Unbox a) =>
SegTree2d s a -> Int -> Int -> Int -> Int -> ST s Int
countST SegTree2d {Int
MVector s a
Vector Int
nSt :: forall s a. SegTree2d s a -> Int
nxSt :: forall s a. SegTree2d s a -> Int
logSt :: forall s a. SegTree2d s a -> Int
sizeSt :: forall s a. SegTree2d s a -> Int
dictXSt :: forall s a. SegTree2d s a -> Vector Int
allYSt :: forall s a. SegTree2d s a -> Vector Int
posSt :: forall s a. SegTree2d s a -> Vector Int
indptrSt :: forall s a. SegTree2d s a -> Vector Int
dataSt :: forall s a. SegTree2d s a -> MVector s a
toLeftSt :: forall s a. SegTree2d s a -> Vector Int
nSt :: Int
nxSt :: Int
logSt :: Int
sizeSt :: Int
dictXSt :: Vector Int
allYSt :: Vector Int
posSt :: Vector Int
indptrSt :: Vector Int
dataSt :: MVector s a
toLeftSt :: Vector Int
..} Int
lx Int
rx Int
ly Int
ry = do
let a0 :: Int
a0 = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
allYSt) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> Maybe Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Maybe Int
lowerBound Vector Int
allYSt Int
ly
let b0 :: Int
b0 = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
allYSt) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> Maybe Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Maybe Int
lowerBound Vector Int
allYSt Int
ry
Int -> Int -> Int -> Int -> Int -> Int -> ST s Int
dfs Int
0 Int
1 Int
0 Int
sizeSt Int
a0 Int
b0
where
!()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
lx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
rx Bool -> Bool -> Bool
&& Int
ly Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
ry) String
"AtCoder.Extra.SegTree2d.countST: given invalid rectangle"
!l0 :: Int
l0 = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
dictXSt) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> Maybe Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Maybe Int
lowerBound Vector Int
dictXSt Int
lx
!r0 :: Int
r0 = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
dictXSt) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int -> Maybe Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Maybe Int
lowerBound Vector Int
dictXSt Int
rx
dfs :: Int -> Int -> Int -> Int -> Int -> Int -> ST s Int
dfs :: Int -> Int -> Int -> Int -> Int -> Int -> ST s Int
dfs (Int
res :: Int) Int
i Int
l Int
r Int
a Int
b
| Int
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
b = Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
res
| Int
r0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l Bool -> Bool -> Bool
|| Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l0 = Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
res
| Int
l0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l Bool -> Bool -> Bool
&& Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r0 = do
Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$! Int
res Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
a
| Bool
otherwise = do
let indptrI :: Int
indptrI = Vector Int
indptrSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
let toLeftI :: Int
toLeftI = Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
indptrI
let la :: Int
la = Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
indptrI Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
a) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
toLeftI
let ra :: Int
ra = Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
la
let lb :: Int
lb = Vector Int
toLeftSt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
indptrI Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
toLeftI
let rb :: Int
rb = Int
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lb
let m :: Int
m = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
Int
res' <- Int -> Int -> Int -> Int -> Int -> Int -> ST s Int
dfs Int
res (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Int
l Int
m Int
la Int
lb
Int -> Int -> Int -> Int -> Int -> Int -> ST s Int
dfs Int
res' (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
m Int
r Int
ra Int
rb