{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE TypeFamilies #-}

{- |
 = Mo's Algotrithm
-}
module Algorithm.Mo where

import Control.Monad.Primitive
import Data.Bits
import qualified Data.Vector.Fusion.Stream.Monadic as MS
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import Data.Word
import Unsafe.Coerce

--

import Data.Vector.Sort.Radix (radixSort64)
import My.Prelude ((..<), (>..))

-- | /O(N * sqrt Q)/
moAlgorithm ::
  (U.Unbox a, PrimMonad m) =>
  -- | index size (N)
  Int ->
  -- | query size (Q)
  Int ->
  -- | add
  (a -> Int -> m a) ->
  -- | delete
  (a -> Int -> m a) ->
  -- | initial value
  a ->
  -- | query [l, r)
  U.Vector MoQuery ->
  m (U.Vector a)
moAlgorithm :: forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Int
-> Int
-> (a -> Int -> m a)
-> (a -> Int -> m a)
-> a
-> Vector MoQuery
-> m (Vector a)
moAlgorithm Int
n Int
q a -> Int -> m a
add a -> Int -> m a
delete a
acc0 Vector MoQuery
qs = do
  MVector (PrimState m) a
result <- Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
UM.unsafeNew Int
q
  (MoState a -> MoQuery -> m (MoState a))
-> MoState a -> Vector MoQuery -> m ()
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m ()
U.foldM'_
    ( \(MoState Int
l Int
r a
acc) (MoQuery Int
ql Int
qr Int
qi) -> do
        (a -> Int -> m a) -> a -> Stream m Int -> m a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Stream m b -> m a
MS.foldM' a -> Int -> m a
add a
acc (Int
l Int -> Int -> Stream m Int
forall (m :: * -> *). Monad m => Int -> Int -> Stream m Int
>.. Int
ql)
          m a -> (a -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (a -> Stream m Int -> m a) -> Stream m Int -> a -> m a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((a -> Int -> m a) -> a -> Stream m Int -> m a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Stream m b -> m a
MS.foldM' a -> Int -> m a
add) (Int
r Int -> Int -> Stream m Int
forall (m :: * -> *). Monad m => Int -> Int -> Stream m Int
..< Int
qr)
          m a -> (a -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (a -> Stream m Int -> m a) -> Stream m Int -> a -> m a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((a -> Int -> m a) -> a -> Stream m Int -> m a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Stream m b -> m a
MS.foldM' a -> Int -> m a
delete) (Int
r Int -> Int -> Stream m Int
forall (m :: * -> *). Monad m => Int -> Int -> Stream m Int
>.. Int
qr)
          m a -> (a -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (a -> Stream m Int -> m a) -> Stream m Int -> a -> m a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((a -> Int -> m a) -> a -> Stream m Int -> m a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Stream m b -> m a
MS.foldM' a -> Int -> m a
delete) (Int
l Int -> Int -> Stream m Int
forall (m :: * -> *). Monad m => Int -> Int -> Stream m Int
..< Int
ql)
          m a -> (a -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
result Int
qi
        Int -> Int -> a -> MoState a
forall a. Int -> Int -> a -> MoState a
MoState Int
ql Int
qr (a -> MoState a) -> m a -> m (MoState a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
result Int
qi
    )
    (Int -> Int -> a -> MoState a
forall a. Int -> Int -> a -> MoState a
MoState Int
0 Int
0 a
acc0)
    (Int -> Vector MoQuery -> Vector MoQuery
moSort (Int -> Int -> Int
moBlockSize Int
n Int
q) Vector MoQuery
qs)
  MVector (PrimState m) a -> m (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector (PrimState m) a
result
{-# INLINE moAlgorithm #-}

-- | @N\/sqrt\ Q@
moBlockSize :: Int -> Int -> Int
moBlockSize :: Int -> Int -> Int
moBlockSize Int
n Int
q = forall a b. (RealFrac a, Integral b) => a -> b
ceiling @Double (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ 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
q)

data MoState a
  = MoState
      !Int
      -- ^ left
      !Int
      -- ^ right
      !a

-- | query for [l,r)
data MoQuery
  = MoQuery
      !Int
      -- ^ left index (inclusive)
      !Int
      -- ^ right index (exclusive)
      !Int
      -- ^ query index

-- | radix sort
moSort ::
  -- | block size
  Int ->
  U.Vector MoQuery ->
  U.Vector MoQuery
moSort :: Int -> Vector MoQuery -> Vector MoQuery
moSort !Int
blockSize Vector MoQuery
qs =
  (Word64 -> MoQuery) -> Vector Word64 -> Vector MoQuery
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
U.map (Vector MoQuery -> Int -> MoQuery
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector MoQuery
qs (Int -> MoQuery) -> (Word64 -> Int) -> Word64 -> MoQuery
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Int
moDecodeQueryIndex)
    (Vector Word64 -> Vector MoQuery)
-> (Vector Word64 -> Vector Word64)
-> Vector Word64
-> Vector MoQuery
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Word64 -> Vector Word64
radixSort64
    (Vector Word64 -> Vector MoQuery)
-> Vector Word64 -> Vector MoQuery
forall a b. (a -> b) -> a -> b
$ (MoQuery -> Word64) -> Vector MoQuery -> Vector Word64
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
U.map (Int -> MoQuery -> Word64
moEncode Int
blockSize) Vector MoQuery
qs
{-# INLINE moSort #-}

-- | query priority
moEncode ::
  -- | block size
  Int ->
  MoQuery ->
  Word64
moEncode :: Int -> MoQuery -> Word64
moEncode !Int
blockSize (MoQuery Int
l Int
r Int
qi) =
  forall a b. a -> b
unsafeCoerce @Int @Word64 (Int -> Word64) -> Int -> Word64
forall a b. (a -> b) -> a -> b
$
    Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
l' Int
40 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
r' Int
20 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
qi
  where
    !l' :: Int
l' = Int -> Int -> Int
forall a. Integral a => a -> a -> a
quot Int
l Int
blockSize
    !r' :: Int
r'
      | Int
l' Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Int
0xfffff Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r
      | Bool
otherwise = Int
r
{-# INLINE moEncode #-}

moDecodeQueryIndex :: Word64 -> Int
moDecodeQueryIndex :: Word64 -> Int
moDecodeQueryIndex = forall a b. a -> b
unsafeCoerce @Word64 @Int (Word64 -> Int) -> (Word64 -> Word64) -> Word64 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0xfffff)
{-# INLINE moDecodeQueryIndex #-}

encodeMoQuery :: MoQuery -> Word64
encodeMoQuery :: MoQuery -> Word64
encodeMoQuery (MoQuery Int
l Int
r Int
i) = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> Int -> Word64
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
l Int
40 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
r Int
20 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
i
{-# INLINE encodeMoQuery #-}

decodeMoQuery :: Word64 -> MoQuery
decodeMoQuery :: Word64 -> MoQuery
decodeMoQuery Word64
w =
  Int -> Int -> Int -> MoQuery
MoQuery
    (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
w) Int
40)
    (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
w) Int
20 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0xfffff)
    (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
w Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0xfffff)
{-# INLINE decodeMoQuery #-}

instance U.IsoUnbox MoQuery Word64 where
  toURepr :: MoQuery -> Word64
toURepr = MoQuery -> Word64
encodeMoQuery
  {-# INLINE toURepr #-}
  fromURepr :: Word64 -> MoQuery
fromURepr = Word64 -> MoQuery
decodeMoQuery
  {-# INLINE fromURepr #-}

newtype instance UM.MVector s MoQuery = MV_MoQuery (UM.MVector s Word64)
newtype instance U.Vector MoQuery = V_MoQuery (U.Vector Word64)
deriving via (MoQuery `U.As` Word64) instance GM.MVector U.MVector MoQuery
deriving via (MoQuery `U.As` Word64) instance G.Vector U.Vector MoQuery
instance U.Unbox MoQuery