{-# LANGUAGE RecordWildCards #-}

-- | A static Wavelet Matrix.
--
-- ==== Notation
-- Let \(S\) be the set of values in your wavelet matrix. We use \(|S|\) to denote the number of
-- distinct values contained within this set \((|S| \lt n)\).
--
-- @since 1.1.0.0
module AtCoder.Extra.WaveletMatrix
  ( -- * Wavelet Matrix
    WaveletMatrix (..),

    -- * Constructors
    build,

    -- * Access (indexing)
    access,

    -- * Rank (count)
    rank,
    rankBetween,

    -- * Selection (finding index)

    -- | ==== __Example__
    -- >>> import AtCoder.Extra.WaveletMatrix qualified as WM
    -- >>> import Data.Vector.Unboxed qualified as VU
    -- >>> let wm = WM.build $ VU.fromList [1,1,2,1,3]
    -- >>> WM.select wm 1
    -- Just 0
    -- >>> WM.selectKth wm 2 1
    -- Just 3
    -- >>> WM.selectIn wm {- [l, r) -} 1 4 {- x -} 1
    -- Just 1
    -- >>> WM.selectKthIn wm {- [l, r) -} 1 4 {- k -} 1 {- x -} 1
    -- Just 3
    select,
    selectKth,
    selectIn,
    selectKthIn,

    -- * Quantile (value-ordered access)
    kthLargestIn,
    ikthLargestIn,
    kthSmallestIn,
    ikthSmallestIn,
    -- unsafeKthLargestIn,
    -- unsafeIKthLargestIn,
    -- unsafeKthSmallestIn,
    -- unsafeIKthSmallestIn,

    -- * Lookup
    lookupLE,
    lookupLT,
    lookupGE,
    lookupGT,

    -- * Conversions
    assocsIn,
    descAssocsIn,
  )
where

import AtCoder.Extra.Bisect
import AtCoder.Extra.WaveletMatrix.Raw qualified as Rwm
import Control.Monad
import Data.Vector.Algorithms.Intro qualified as VAI
import Data.Vector.Generic qualified as VG
import Data.Vector.Unboxed qualified as VU
import GHC.Stack (HasCallStack)

-- | A static Wavelet Matrix.
--
-- @since 1.3.0.0
data WaveletMatrix = WaveletMatrix
  { -- | The internal wavelet matrix, where index compression is not automatically performed.
    --
    -- @since 1.3.0.0
    WaveletMatrix -> RawWaveletMatrix
rawWm :: !Rwm.RawWaveletMatrix,
    -- | Index compression dictionary.
    --
    -- @since 1.3.0.0
    WaveletMatrix -> Vector Int
yDictWm :: !(VU.Vector Int)
  }

-- | \(O(n \log n)\) Creates a `WaveletMatrix` from an array \(a\).
--
-- @since 1.1.0.0
{-# INLINE build #-}
build :: VU.Vector Int -> WaveletMatrix
build :: Vector Int -> WaveletMatrix
build Vector Int
ys =
  let !yDictWm :: Vector Int
yDictWm = 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 (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 Comparison Int
forall a. Ord a => a -> a -> Ordering
compare) Vector Int
ys
      !ys' :: Vector Int
ys' = (Int -> Int) -> Vector Int -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Int
lowerBound Vector Int
yDictWm) Vector Int
ys
      !rawWm :: RawWaveletMatrix
rawWm = HasCallStack => Int -> Vector Int -> RawWaveletMatrix
Int -> Vector Int -> RawWaveletMatrix
Rwm.build (Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector Int
ys) Vector Int
ys'
   in WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
yDictWm :: Vector Int
rawWm :: RawWaveletMatrix
..}

-- | \(O(\log |S|)\) Returns \(a[k]\) or `Nothing` if the index is out of the bounds. Try to use the
-- original array if you can.
--
-- @since 1.1.0.0
{-# INLINE access #-}
access ::  WaveletMatrix -> Int -> Maybe Int
access :: WaveletMatrix -> Int -> Maybe Int
access WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
i = (Vector Int
yDictWm VG.!) (Int -> Int) -> Maybe Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RawWaveletMatrix -> Int -> Maybe Int
Rwm.access RawWaveletMatrix
rawWm Int
i

-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\).
--
-- @since 1.1.0.0
{-# INLINE rank #-}
rank ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(y\)
  Int ->
  -- | The number of \(y\) in \([l, r)\).
  Int
rank :: WaveletMatrix -> Int -> Int -> Int -> Int
rank WaveletMatrix
wm Int
l Int
r Int
y = WaveletMatrix -> Int -> Int -> Int -> Int -> Int
rankBetween WaveletMatrix
wm Int
l Int
r Int
y (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\).
--
-- @since 1.1.0.0
{-# INLINEABLE rankBetween #-}
rankBetween ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(y_1\)
  Int ->
  -- | \(y_2\)
  Int ->
  -- | The number of \(y\) in \([l, r) \times [y_1, y_2)\).
  Int
rankBetween :: WaveletMatrix -> Int -> Int -> Int -> Int -> Int
rankBetween WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
l Int
r Int
y1 Int
y2
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l Bool -> Bool -> Bool
&& Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r Bool -> Bool -> Bool
&& Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n = Int
0
  | Int
y1' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
y2' = Int
0
  | Bool
otherwise = RawWaveletMatrix -> Int -> Int -> Int -> Int -> Int
Rwm.rankBetween RawWaveletMatrix
rawWm Int
l Int
r Int
y1' Int
y2'
  where
    -- Handles the case @yl@ or  @yr@ is not in the dict
    n :: Int
n = RawWaveletMatrix -> Int
Rwm.lengthRwm RawWaveletMatrix
rawWm
    y1' :: Int
y1' = Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Int
lowerBound Vector Int
yDictWm Int
y1
    y2' :: Int
y2' = Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Int
lowerBound Vector Int
yDictWm Int
y2

-- | \(O(\log |S|)\) Returns the index of the first \(y\) in \(a\), or `Nothing` if \(y\) is
-- not found.
--
-- @since 1.1.0.0
{-# INLINE select #-}
select :: WaveletMatrix -> Int -> Maybe Int
select :: WaveletMatrix -> Int -> Maybe Int
select WaveletMatrix
wm = WaveletMatrix -> Int -> Int -> Maybe Int
selectKth WaveletMatrix
wm Int
0

-- | \(O(\log |S|)\) Returns the index of the \(k\)-th occurrence (0-based) of \(y\), or `Nothing`
-- if no such occurrence exists.
--
-- @since 1.1.0.0
{-# INLINEABLE selectKth #-}
selectKth ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(k\)
  Int ->
  -- | \(y\)
  Int ->
  -- | The index of \(k\)-th \(y\)
  Maybe Int
selectKth :: WaveletMatrix -> Int -> Int -> Maybe Int
selectKth WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
k Int
y = do
  let !i :: Int
i = Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Int
lowerBound Vector Int
yDictWm Int
y
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector Int
yDictWm
  -- TODO: we don't need such an explicit branch?
  let !y' :: Int
y' = Vector Int
yDictWm Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Int
y' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y
  RawWaveletMatrix -> Int -> Int -> Maybe Int
Rwm.selectKth RawWaveletMatrix
rawWm Int
k Int
i

-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the first occurrence
-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists.
--
-- @since 1.1.0.0
{-# INLINE selectIn #-}
selectIn ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(y\)
  Int ->
  -- | The index of the first \(y\) in \([l, r)\).
  Maybe Int
selectIn :: WaveletMatrix -> Int -> Int -> Int -> Maybe Int
selectIn WaveletMatrix
wm Int
l Int
r = HasCallStack =>
WaveletMatrix -> Int -> Int -> Int -> Int -> Maybe Int
WaveletMatrix -> Int -> Int -> Int -> Int -> Maybe Int
selectKthIn WaveletMatrix
wm Int
l Int
r Int
0

-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the \(k\)-th occurrence
-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists.
--
-- @since 1.1.0.0
{-# INLINEABLE selectKthIn #-}
selectKthIn ::
  (HasCallStack) =>
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(k\)
  Int ->
  -- | \(y\)
  Int ->
  -- | The index of the \(k\)-th \(y\) in \([l, r)\).
  Maybe Int
selectKthIn :: HasCallStack =>
WaveletMatrix -> Int -> Int -> Int -> Int -> Maybe Int
selectKthIn WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
l Int
r Int
k Int
y = do
  let !i :: Int
i = Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Ord a) =>
v a -> a -> Int
lowerBound Vector Int
yDictWm Int
y
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector Int
yDictWm
  -- TODO: we don't need such an explicit branch?
  let !y' :: Int
y' = Vector Int
yDictWm Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Int
y' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y
  RawWaveletMatrix -> Int -> Int -> Int -> Int -> Maybe Int
Rwm.selectKthIn RawWaveletMatrix
rawWm Int
l Int
r Int
k Int
i

-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns the index of the \(k\)-th (0-based)
-- largest value in it. Note that duplicated values are treated as distinct occurrences.
--
-- @since 1.1.0.0
{-# INLINEABLE kthLargestIn #-}
kthLargestIn ::
  (HasCallStack) =>
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(k\)
  Int ->
  -- | \(k\)-th largest \(y\) in \([l, r)\)
  Maybe Int
kthLargestIn :: HasCallStack => WaveletMatrix -> Int -> Int -> Int -> Maybe Int
kthLargestIn WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
l Int
r Int
k
  | Just !Int
y <- RawWaveletMatrix -> Int -> Int -> Int -> Maybe Int
Rwm.kthLargestIn RawWaveletMatrix
rawWm Int
l Int
r Int
k = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Vector Int
yDictWm Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
y
  | Bool
otherwise = Maybe Int
forall a. Maybe a
Nothing

-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns both the index and the value of the
-- \(k\)-th (0-based) largest value in it. Note that duplicated values are treated as distinct
-- occurrences.
--
-- @since 1.1.0.0
{-# INLINEABLE ikthLargestIn #-}
ikthLargestIn ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(k\)
  Int ->
  -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\)
  Maybe (Int, Int)
ikthLargestIn :: WaveletMatrix -> Int -> Int -> Int -> Maybe (Int, Int)
ikthLargestIn WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
l Int
r Int
k
  | Just (!Int
i, !Int
y) <- RawWaveletMatrix -> Int -> Int -> Int -> Maybe (Int, Int)
Rwm.ikthLargestIn RawWaveletMatrix
rawWm Int
l Int
r Int
k = (Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
i, Vector Int
yDictWm Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
y)
  | Bool
otherwise = Maybe (Int, Int)
forall a. Maybe a
Nothing

-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns the index of the \(k\)-th (0-based)
-- smallest value in it. Note that duplicated values are treated as distinct occurrences.
--
-- @since 1.1.0.0
{-# INLINEABLE kthSmallestIn #-}
kthSmallestIn ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(k\)
  Int ->
  -- | \(k\)-th largest \(y\) in \([l, r)\)
  Maybe Int
kthSmallestIn :: WaveletMatrix -> Int -> Int -> Int -> Maybe Int
kthSmallestIn WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
l Int
r Int
k
  | Just !Int
y <- RawWaveletMatrix -> Int -> Int -> Int -> Maybe Int
Rwm.kthSmallestIn RawWaveletMatrix
rawWm Int
l Int
r Int
k = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Vector Int
yDictWm Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
y
  | Bool
otherwise = Maybe Int
forall a. Maybe a
Nothing

-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns both the index and the value of the
-- \(k\)-th (0-based) smallest value in it. Note that duplicated values are treated as distinct
-- occurrences.
--
-- @since 1.1.0.0
{-# INLINEABLE ikthSmallestIn #-}
ikthSmallestIn ::
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(k\)
  Int ->
  -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\)
  Maybe (Int, Int)
ikthSmallestIn :: WaveletMatrix -> Int -> Int -> Int -> Maybe (Int, Int)
ikthSmallestIn WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
l Int
r Int
k
  | Just (!Int
i, !Int
y) <- RawWaveletMatrix -> Int -> Int -> Int -> Maybe (Int, Int)
Rwm.ikthSmallestIn RawWaveletMatrix
rawWm Int
l Int
r Int
k = (Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
i, Vector Int
yDictWm Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
y)
  | Bool
otherwise = Maybe (Int, Int)
forall a. Maybe a
Nothing

-- | \(O(\log |S|)\)
--
-- @since 1.1.0.0
{-# INLINE unsafeKthSmallestIn #-}
unsafeKthSmallestIn :: WaveletMatrix -> Int -> Int -> Int -> Int
unsafeKthSmallestIn :: WaveletMatrix -> Int -> Int -> Int -> Int
unsafeKthSmallestIn WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
l Int
r Int
k =
  Vector Int
yDictWm Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! RawWaveletMatrix -> Int -> Int -> Int -> Int
Rwm.unsafeKthSmallestIn RawWaveletMatrix
rawWm Int
l Int
r Int
k

-- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0]\).
--
-- @since 1.1.0.0
{-# INLINEABLE lookupLE #-}
lookupLE ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(y_0\)
  Int ->
  -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0]\)
  Maybe Int
lookupLE :: WaveletMatrix -> Int -> Int -> Int -> Maybe Int
lookupLE WaveletMatrix
wm Int
l Int
r Int
y0
  | Int
r' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l' = Maybe Int
forall a. Maybe a
Nothing
  | Int
rank_ Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Maybe Int
forall a. Maybe a
Nothing
  | Bool
otherwise = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ WaveletMatrix -> Int -> Int -> Int -> Int
unsafeKthSmallestIn WaveletMatrix
wm Int
l' Int
r' (Int
rank_ Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    -- clamp
    l' :: Int
l' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
l
    r' :: Int
r' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (RawWaveletMatrix -> Int
Rwm.lengthRwm (WaveletMatrix -> RawWaveletMatrix
rawWm WaveletMatrix
wm)) Int
r
    rank_ :: Int
rank_ = WaveletMatrix -> Int -> Int -> Int -> Int -> Int
rankBetween WaveletMatrix
wm Int
l' Int
r' Int
forall a. Bounded a => a
minBound (Int
y0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0)\).
--
-- @since 1.1.0.0
{-# INLINE lookupLT #-}
lookupLT ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(y_0\)
  Int ->
  -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0)\)
  Maybe Int
lookupLT :: WaveletMatrix -> Int -> Int -> Int -> Maybe Int
lookupLT WaveletMatrix
wm Int
l Int
r Int
y0 = WaveletMatrix -> Int -> Int -> Int -> Maybe Int
lookupLE WaveletMatrix
wm Int
l Int
r (Int
y0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times [y_0, \infty)\).
--
-- @since 1.1.0.0
{-# INLINEABLE lookupGE #-}
lookupGE ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(y_0\)
  Int ->
  -- | Minimum \(y\) in \([l, r) \times [y_0, \infty)\).
  Maybe Int
lookupGE :: WaveletMatrix -> Int -> Int -> Int -> Maybe Int
lookupGE WaveletMatrix
wm Int
l Int
r Int
y0
  | Int
r' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l' = Maybe Int
forall a. Maybe a
Nothing
  | Int
rank_ Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l = Maybe Int
forall a. Maybe a
Nothing
  | Bool
otherwise = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ WaveletMatrix -> Int -> Int -> Int -> Int
unsafeKthSmallestIn WaveletMatrix
wm Int
l Int
r Int
rank_
  where
    -- clamp
    l' :: Int
l' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
l
    r' :: Int
r' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (RawWaveletMatrix -> Int
Rwm.lengthRwm (WaveletMatrix -> RawWaveletMatrix
rawWm WaveletMatrix
wm)) Int
r
    rank_ :: Int
rank_ = WaveletMatrix -> Int -> Int -> Int -> Int -> Int
rankBetween WaveletMatrix
wm Int
l' Int
r' Int
forall a. Bounded a => a
minBound Int
y0

-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\).
--
-- @since 1.1.0.0
{-# INLINE lookupGT #-}
lookupGT ::
  -- | A wavelet matrix
  WaveletMatrix ->
  -- | \(l\)
  Int ->
  -- | \(r\)
  Int ->
  -- | \(y_0\)
  Int ->
  -- | Minimum \(y\) in \([l, r) \times (y_0, \infty)\)
  Maybe Int
lookupGT :: WaveletMatrix -> Int -> Int -> Int -> Maybe Int
lookupGT WaveletMatrix
wm Int
l Int
r Int
y0 = WaveletMatrix -> Int -> Int -> Int -> Maybe Int
lookupGE WaveletMatrix
wm Int
l Int
r (Int
y0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in an interval \([l, r)\) in
-- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small.
--
-- @since 1.1.0.0
{-# INLINE assocsIn #-}
assocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)]
assocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)]
assocsIn WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
l Int
r = RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)]
Rwm.assocsWith RawWaveletMatrix
rawWm Int
l Int
r (Vector Int
yDictWm VG.!)

-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in an interval \([l, r)\) in
-- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small.
--
-- @since 1.1.0.0
{-# INLINE descAssocsIn #-}
descAssocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)]
descAssocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)]
descAssocsIn WaveletMatrix {Vector Int
RawWaveletMatrix
rawWm :: WaveletMatrix -> RawWaveletMatrix
yDictWm :: WaveletMatrix -> Vector Int
rawWm :: RawWaveletMatrix
yDictWm :: Vector Int
..} Int
l Int
r = RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)]
Rwm.descAssocsInWith RawWaveletMatrix
rawWm Int
l Int
r (Vector Int
yDictWm VG.!)