{-|
== Mo's algorithm

Useful for answering certain range queries on a static sequence. The algorithm works by maintaining
some state for a current range, and updating the state by adding or removing elements one by one.
An important parameter for the algorithm is the block size. If the range size is n, number of
queries is q, block size is b, state update takes time f(n), getting the answer for the current
state takes g(n), then the algorithm answers all queries in
    O(q log q + bq * f(n) + n^2/b * f(n) + q * g(n))
A suitable value of block size minimizes run time.
For example, when q ~ n, f(n) = O(1), g(n) = O(1), setting b = sqrt(n) gives run time O(n sqrt(n)).
For f(n) = O(log n), a better choice of b is sqrt(n/log n) giving run time O(n sqrt(n log n)). In
practice, it works to experiment a bit and choose a good fixed value.

Sources:

* https://cp-algorithms.com/data_structures/sqrt_decomposition.html

-}

module Mo
    ( MoQuery(..)
    , Tag
    , runMo
    , sqrtSize
    ) where

import Control.DeepSeq
import Data.List

type Tag = Int
data MoQuery = MoQuery { MoQuery -> Int
ql_ :: !Int, MoQuery -> Int
qr_ :: !Int, MoQuery -> Int
qtag_ :: !Tag } deriving Int -> MoQuery -> ShowS
[MoQuery] -> ShowS
MoQuery -> String
(Int -> MoQuery -> ShowS)
-> (MoQuery -> String) -> ([MoQuery] -> ShowS) -> Show MoQuery
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MoQuery] -> ShowS
$cshowList :: [MoQuery] -> ShowS
show :: MoQuery -> String
$cshow :: MoQuery -> String
showsPrec :: Int -> MoQuery -> ShowS
$cshowsPrec :: Int -> MoQuery -> ShowS
Show

-- | Run Mo's algorithm given the block size, state update and answer functions, and queries. See above
-- for time complexity.
runMo :: Monad m => Int -> (Int -> m ()) -> (Int -> m ()) -> m a -> [MoQuery] -> m [(Tag, a)]
runMo :: Int
-> (Int -> m ())
-> (Int -> m ())
-> m a
-> [MoQuery]
-> m [(Int, a)]
runMo Int
_     Int -> m ()
_   Int -> m ()
_   m a
_   []   = [(Int, a)] -> m [(Int, a)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
runMo Int
bsize Int -> m ()
add Int -> m ()
rem m a
ans [MoQuery]
qrys = [MoQuery] -> Int -> Int -> [(Int, a)] -> m [(Int, a)]
go [MoQuery]
qrys' Int
start (Int
startInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [] where
    cmp :: MoQuery -> MoQuery -> Ordering
cmp (MoQuery Int
l1 Int
r1 Int
_) (MoQuery Int
l2 Int
r2 Int
_) = Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
b1 Int
b2 Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Ordering
rc where
        (Int
b1, Int
b2) = (Int
l1 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
bsize, Int
l2 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
bsize)
        rc :: Ordering
rc = if Int -> Bool
forall a. Integral a => a -> Bool
even Int
b1 then Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
r1 Int
r2 else Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
r2 Int
r1
    qrys' :: [MoQuery]
qrys' = (MoQuery -> MoQuery -> Ordering) -> [MoQuery] -> [MoQuery]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy MoQuery -> MoQuery -> Ordering
cmp [MoQuery]
qrys
    MoQuery Int
start Int
_ Int
_ = [MoQuery] -> MoQuery
forall a. [a] -> a
head [MoQuery]
qrys'
    go :: [MoQuery] -> Int -> Int -> [(Int, a)] -> m [(Int, a)]
go [] Int
_ Int
_ [(Int, a)]
acc = [(Int, a)] -> m [(Int, a)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Int, a)]
acc
    go ((MoQuery Int
ql Int
qr Int
qtag):[MoQuery]
qrys) Int
l Int
r [(Int, a)]
acc = do
        (Int -> m ()) -> [Int] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Int -> m ()
add ([Int] -> m ()) -> [Int] -> m ()
forall a b. (a -> b) -> a -> b
$ [Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1, Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2 .. Int
ql] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1 .. Int
qr]
        (Int -> m ()) -> [Int] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Int -> m ()
rem ([Int] -> m ()) -> [Int] -> m ()
forall a b. (a -> b) -> a -> b
$ [Int
l .. Int
qlInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
r, Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 .. Int
qrInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1]
        a
x <- m a
ans
        [MoQuery] -> Int -> Int -> [(Int, a)] -> m [(Int, a)]
go [MoQuery]
qrys Int
ql Int
qr ((Int
qtag, a
x)(Int, a) -> [(Int, a)] -> [(Int, a)]
forall a. a -> [a] -> [a]
:[(Int, a)]
acc)

-- | Square root of an Int, rounded to an Int. O(1).
sqrtSize :: Int -> Int
sqrtSize :: Int -> Int
sqrtSize Int
n = Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
round (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n :: Double)

--------------------------------------------------------------------------------
-- For tests

-- Allows specialization across modules
{-# INLINABLE runMo #-}

instance NFData MoQuery where
    rnf :: MoQuery -> ()
rnf = MoQuery -> ()
forall a. a -> ()
rwhnf