-- | Mo's algorithm for handling \([l, r)\) offline queries in \(O((n + q) \sqrt n f)\) time
-- complecity, where \(n\) is the length of index, \(q\) is the number of queries and \(f\) is the
-- time for processing element addition or deletion. Due to the high time complexity, it is
-- recommended to choose an efficient data structure such as Fenwick Tree for query processing.
--
-- @since 1.2.5.0
module AtCoder.Extra.Mo
  ( run,
    sortIndices,
    process,
  )
where

import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad)
import Data.Bool (bool)
import Data.Foldable (for_)
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)

-- | \(O((n + q) \sqrt n)\) Runs Mo's algorithm. Internally it's a call of `sortIndices` and
-- `process`.
--
-- @since 1.2.5.0
{-# INLINE run #-}
run ::
  (HasCallStack, PrimMonad m, VU.Unbox a) =>
  -- | Defines index bounds \([0, n)\).
  Int ->
  -- | Query intervals \([l, r)\).
  VU.Vector (Int, Int) ->
  -- | Called on adding left index \(l\).
  (Int -> m ()) ->
  -- | Called on adding left index \(r\).
  (Int -> m ()) ->
  -- | Called on deleting left index \(l\).
  (Int -> m ()) ->
  -- | Called on deleting right index \(r\).
  (Int -> m ()) ->
  -- | Returns result for query index \(i\).
  (Int -> m a) ->
  -- | Result for each query.
  m (VU.Vector a)
run :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Int
-> Vector (Int, Int)
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m a)
-> m (Vector a)
run Int
n !Vector (Int, Int)
lrs !Int -> m ()
addL !Int -> m ()
addR !Int -> m ()
delL !Int -> m ()
delR !Int -> m a
query = do
  let !is :: Vector Int
is = HasCallStack => Int -> Vector (Int, Int) -> Vector Int
Int -> Vector (Int, Int) -> Vector Int
sortIndices Int
n Vector (Int, Int)
lrs
  Vector Int
-> Vector (Int, Int)
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m a)
-> m (Vector a)
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Vector Int
-> Vector (Int, Int)
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m a)
-> m (Vector a)
process Vector Int
is Vector (Int, Int)
lrs Int -> m ()
addL Int -> m ()
addR Int -> m ()
delL Int -> m ()
delR Int -> m a
query

-- | \(O(n (\log n))\) Sorts indices of \([l, r)\) queries in an efficient order for processing.
--
-- @since 1.2.5.0
{-# INLINEABLE sortIndices #-}
sortIndices ::
  (HasCallStack) =>
  -- | Defines index bounds \([0, n)\).
  Int ->
  -- | Query intervals \([l, r)\).
  VU.Vector (Int, Int) ->
  -- | Sorted indices to the query intervals.
  VU.Vector Int
sortIndices :: HasCallStack => Int -> Vector (Int, Int) -> Vector Int
sortIndices Int
n !Vector (Int, Int)
lrs
  | Vector (Int, Int) -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector (Int, Int)
lrs = Vector Int
forall a. Unbox a => Vector a
VU.empty
  | Bool
otherwise = (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
      let !q :: Int
q = Vector (Int, Int) -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector (Int, Int)
lrs
      let !Int
blockLen :: Int = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double -> Double
forall a. Floating a => a -> a
sqrt Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
q)) :: Double)
      MVector (PrimState (ST s)) Int
is <- Int -> (Int -> Int) -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> (Int -> a) -> m (MVector (PrimState m) a)
VUM.generate Int
q Int -> Int
forall a. a -> a
id

      -- sort by block index then right:
      Comparison Int -> MVector (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
VAI.sortBy
        ( \Int
i1 Int
i2 -> do
            let (!Int
l1, !Int
r1) = Vector (Int, Int)
lrs Vector (Int, Int) -> Int -> (Int, Int)
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i1
                (!Int
l2, !Int
r2) = Vector (Int, Int)
lrs Vector (Int, Int) -> Int -> (Int, Int)
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i2
                !b1 :: Int
b1 = Int
l1 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
blockLen
                !b2 :: Int
b2 = Int
l2 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
blockLen
                !res :: Ordering
res = Comparison Int
forall a. Ord a => a -> a -> Ordering
compare Int
b1 Int
b2 Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Ordering -> Ordering -> Bool -> Ordering
forall a. a -> a -> Bool -> a
bool (Comparison Int
forall a. Ord a => a -> a -> Ordering
compare Int
r2 Int
r1) (Comparison Int
forall a. Ord a => a -> a -> Ordering
compare Int
r1 Int
r2) (Int -> Bool
forall a. Integral a => a -> Bool
even Int
b1)
             in Ordering
res
        )
        MVector (PrimState (ST s)) Int
is

      -- The following trick doesn't seem to make it faster though?

      let -- {-# INLINE cost #-}
          cost :: Int -> Int -> ST s Int
cost Int
i1 Int
i2 = do
            (!Int
l1, !Int
r1) <- (Vector (Int, Int)
lrs VG.!) (Int -> (Int, Int)) -> ST s Int -> ST s (Int, Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
is Int
i1
            (!Int
l2, !Int
r2) <- (Vector (Int, Int)
lrs VG.!) (Int -> (Int, Int)) -> ST s Int -> ST s (Int, Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
is Int
i2
            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 -> Int
forall a. Num a => a -> a
abs (Int
l1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l2) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Num a => a -> a
abs (Int
r1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r2)

      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
q Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
6) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
6] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
k -> do
          do
            Int
c1 <- Int -> Int -> ST s Int
cost Int
k (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
            Int
c2 <- Int -> Int -> ST s Int
cost (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3)
            Int
c3 <- Int -> Int -> ST s Int
cost Int
k (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Int
c4 <- Int -> Int -> ST s Int
cost (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3)
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
c1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
c3 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c4) (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 -> Int -> m ()
VGM.swap MVector (PrimState (ST s)) Int
is (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
          do
            Int
c1 <- Int -> Int -> ST s Int
cost Int
k (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3)
            Int
c2 <- Int -> Int -> ST s Int
cost (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4)
            Int
c3 <- Int -> Int -> ST s Int
cost Int
k (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Int
c4 <- Int -> Int -> ST s Int
cost (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4)
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
c1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
c3 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c4) (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 -> Int -> m ()
VGM.swap MVector (PrimState (ST s)) Int
is (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3)

      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
is

-- | \(O((n + q) \sqrt n)\) Processes \([l, r)\) interval queries. User would usually use `run`
-- instead.
--
-- @since 1.2.5.0
{-# INLINEABLE process #-}
process ::
  (HasCallStack, PrimMonad m, VU.Unbox a) =>
  -- | Sorted indices to query intervals \([l, r)\).
  VU.Vector Int ->
  -- | Query intervals \([l, r)\).
  VU.Vector (Int, Int) ->
  -- | Called on adding left index \(l\).
  (Int -> m ()) ->
  -- | Called on adding right index \(r\).
  (Int -> m ()) ->
  -- | Called on deleting left index \(l\).
  (Int -> m ()) ->
  -- | Called on deleting right index \(r\).
  (Int -> m ()) ->
  -- | Returns result for query index \(i\).
  (Int -> m a) ->
  -- | Result for each query.
  m (VU.Vector a)
process :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Vector Int
-> Vector (Int, Int)
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m ())
-> (Int -> m a)
-> m (Vector a)
process !Vector Int
is !Vector (Int, Int)
lrs !Int -> m ()
addL !Int -> m ()
addR !Int -> m ()
delL !Int -> m ()
delR !Int -> m a
query = do
  let !q :: Int
q = Vector (Int, Int) -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector (Int, Int)
lrs
  !MVector (PrimState m) a
result <- Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
q

  ((Int, Int) -> Int -> m (Int, Int))
-> (Int, Int) -> Vector Int -> m ()
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m ()
VU.foldM'_
    ( \(!Int
l0, !Int
r0) Int
i -> do
        let (!Int
l, !Int
r) = Vector (Int, Int)
lrs Vector (Int, Int) -> Int -> (Int, Int)
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
        [Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
l] Int -> m ()
addL
        [Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
r0, Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] Int -> m ()
addR
        [Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
l0, Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 .. Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] Int -> m ()
delL
        [Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
r] Int -> m ()
delR
        MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector (PrimState m) a
result Int
i (a -> m ()) -> m a -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> m a
query Int
i
        (Int, Int) -> m (Int, Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
l, Int
r)
    )
    (Int
0 :: Int, Int
0 :: Int)
    Vector Int
is

  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
result