{-# LANGUAGE RecordWildCards #-}

-- original implementation:
-- <https://github.com/maspypy/library/blob/main/ds/fastset.hpp>

-- | A dense, fast `Int` set implemented as a 64-ary tree that covers an interval \([0, n)\).
--
-- ==== __Example__
-- Create an `IntSet` with capacity \(10\):
--
-- >>> import AtCoder.Extra.IntSet qualified as IS
-- >>> is <- IS.new @_ 10
--
-- `insert`, `delete` and other functions are available:
--
-- >>> IS.insert is 0
-- >>> IS.insert is 9
-- >>> IS.member is 9
-- True
--
-- >>> IS.delete is 0
-- True
--
-- >>> IS.size is
-- 1
--
-- >>> IS.member is 1
-- False
--
-- >>> IS.lookupGT is 5
-- Just 9
--
-- @since 1.1.0.0
module AtCoder.Extra.IntSet
  ( -- * IntSet
    IntSet,

    -- * Constructors
    new,
    build,

    -- * Metadata
    capacity,
    size,
    null,

    -- * Lookups
    member,
    notMember,

    -- ** Compartive lookups
    lookupGE,
    lookupGT,
    lookupLE,
    lookupLT,

    -- ** Max/Min lookups
    lookupMin,
    lookupMax,

    -- * Modifications

    -- ** Insertions
    insert,

    -- ** Deletions
    delete,
    delete_,
    deleteMin,
    deleteMax,

    -- * Conversions
    keys,
  )
where

import AtCoder.Internal.Assert qualified as ACIA
import Control.Monad (unless, void)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Bifunctor (bimap)
import Data.Bits
  ( Bits (clearBit, setBit, testBit),
    FiniteBits (countLeadingZeros, countTrailingZeros),
    (.<<.),
    (.>>.),
  )
import Data.Maybe (fromJust)
import Data.Vector qualified as V
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 (null)

-- | \(O(1)\) Retrieves the most significant bit.
--
-- >>> msbOf 0
-- -1
--
-- >>> msbOf maxBound
-- 62
--
-- >>> msbOf $ 4 + 2 + 1
-- 2
{-# INLINE msbOf #-}
msbOf :: Int -> Int
msbOf :: Int -> Int
msbOf !Int
x = Int
63 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Int
x

-- | \(O(1)\) Retrieves the least significant bit.
--
-- >>> lsbOf 0
-- -1
--
-- >>> lsbOf maxBound
-- 0
--
-- >>> lsbOf $ 4 + 2 + 1
-- 0
{-# INLINE lsbOf #-}
lsbOf :: Int -> Int
lsbOf :: Int -> Int
lsbOf Int
0 = -Int
1
lsbOf Int
x = Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
x

{-# INLINE wordSize #-}
wordSize :: Int
wordSize :: Int
wordSize = Int
64

-- | A dense, fast `Int` set implemented as a 64-ary tree that covers an interval \([0, n)\).
--
-- @since 1.1.0.0
data IntSet s = IntSet
  { -- | Maximum number of elements.
    forall s. IntSet s -> Int
capacityIS :: {-# UNPACK #-} !Int,
    -- | The number of elements.
    forall s. IntSet s -> MVector s Int
sizeIS :: !(VUM.MVector s Int),
    -- | Segments.
    forall s. IntSet s -> Vector (MVector s Int)
vecIS :: !(V.Vector (VUM.MVector s Int))
  }

-- | \(O(n)\) Creates an `IntSet` for the interval \([0, n)\).
--
-- @since 1.1.0.0
{-# INLINE new #-}
new :: (PrimMonad m) => Int -> m (IntSet (PrimState m))
new :: forall (m :: * -> *).
PrimMonad m =>
Int -> m (IntSet (PrimState m))
new Int
capacityIS = do
  Vector (MVector (PrimState m) Int)
vecIS <-
    Int
-> (Int -> m (MVector (PrimState m) Int, Int))
-> Int
-> m (Vector (MVector (PrimState m) Int))
forall (m :: * -> *) b a.
Monad m =>
Int -> (b -> m (a, b)) -> b -> m (Vector a)
V.unfoldrExactNM
      (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
logSize)
      ( \Int
len -> do
          let !len' :: Int
len' = (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
wordSize
          (,Int
len') (MVector (PrimState m) Int -> (MVector (PrimState m) Int, Int))
-> m (MVector (PrimState m) Int)
-> m (MVector (PrimState m) Int, Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
len' Int
0
      )
      Int
capacityIS
  MVector (PrimState m) Int
sizeIS <- Int -> Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 (Int
0 :: Int)
  IntSet (PrimState m) -> m (IntSet (PrimState m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntSet {Int
MVector (PrimState m) Int
Vector (MVector (PrimState m) Int)
capacityIS :: Int
sizeIS :: MVector (PrimState m) Int
vecIS :: Vector (MVector (PrimState m) Int)
capacityIS :: Int
vecIS :: Vector (MVector (PrimState m) Int)
sizeIS :: MVector (PrimState m) Int
..}
  where
    (!Int
_, !Int
logSize) =
      ((Int, Int) -> Bool)
-> ((Int, Int) -> (Int, Int)) -> (Int, Int) -> (Int, Int)
forall a. (a -> Bool) -> (a -> a) -> a -> a
until
        ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1) (Int -> Bool) -> ((Int, Int) -> Int) -> (Int, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> Int
forall a b. (a, b) -> a
fst)
        ((Int -> Int) -> (Int -> Int) -> (Int, Int) -> (Int, Int)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap ((Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
wordSize) (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        (Int
capacityIS, Int
0)

-- | \(O(n + m \log n)\) Creates an `IntSet` for the interval \([0, n)\) with initial values.
--
-- @since 1.1.0.0
{-# INLINE build #-}
build :: (PrimMonad m) => Int -> VU.Vector Int -> m (IntSet (PrimState m))
build :: forall (m :: * -> *).
PrimMonad m =>
Int -> Vector Int -> m (IntSet (PrimState m))
build Int
n Vector Int
vs = do
  IntSet (PrimState m)
set <- Int -> m (IntSet (PrimState m))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (IntSet (PrimState m))
new Int
n
  Vector Int -> (Int -> m ()) -> m ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector Int
vs (IntSet (PrimState m) -> Int -> m ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
IntSet (PrimState m) -> Int -> m ()
insert IntSet (PrimState m)
set)
  IntSet (PrimState m) -> m (IntSet (PrimState m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntSet (PrimState m)
set

-- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is covered by the set.
--
-- @since 1.1.0.0
{-# INLINE capacity #-}
capacity :: IntSet s -> Int
capacity :: forall s. IntSet s -> Int
capacity = IntSet s -> Int
forall s. IntSet s -> Int
capacityIS

-- | \(O(1)\) Returns the number of elements in the set.
--
-- @since 1.1.0.0
{-# INLINE size #-}
size :: (PrimMonad m) => IntSet (PrimState m) -> m Int
size :: forall (m :: * -> *). PrimMonad m => IntSet (PrimState m) -> m Int
size = (MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
`VUM.unsafeRead` Int
0) (MVector (PrimState m) Int -> m Int)
-> (IntSet (PrimState m) -> MVector (PrimState m) Int)
-> IntSet (PrimState m)
-> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet (PrimState m) -> MVector (PrimState m) Int
forall s. IntSet s -> MVector s Int
sizeIS

-- | \(O(1)\) Returns whether the set is empty.
--
-- @since 1.1.0.0
{-# INLINE null #-}
null :: (PrimMonad m) => IntSet (PrimState m) -> m Bool
null :: forall (m :: * -> *). PrimMonad m => IntSet (PrimState m) -> m Bool
null = ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) <$>) (m Int -> m Bool)
-> (IntSet (PrimState m) -> m Int)
-> IntSet (PrimState m)
-> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet (PrimState m) -> m Int
forall (m :: * -> *). PrimMonad m => IntSet (PrimState m) -> m Int
size

-- | \(O(\log n)\) Tests whether \(k\) is in the set.
--
-- @since 1.1.0.0
{-# INLINE member #-}
member :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool
member :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
member IntSet {Int
MVector (PrimState m) Int
Vector (MVector (PrimState m) Int)
capacityIS :: forall s. IntSet s -> Int
sizeIS :: forall s. IntSet s -> MVector s Int
vecIS :: forall s. IntSet s -> Vector (MVector s Int)
capacityIS :: Int
sizeIS :: MVector (PrimState m) Int
vecIS :: Vector (MVector (PrimState m) Int)
..} Int
k
  | HasCallStack => Int -> Int -> Bool
Int -> Int -> Bool
ACIA.testIndex Int
k Int
capacityIS = do
      let (!Int
q, !Int
r) = Int
k Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
wordSize
      (Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
`testBit` Int
r) (Int -> Bool) -> m Int -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead (Vector (MVector (PrimState m) Int) -> MVector (PrimState m) Int
forall (v :: * -> *) a. Vector v a => v a -> a
VG.unsafeHead Vector (MVector (PrimState m) Int)
vecIS) Int
q
  | Bool
otherwise = Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

-- | \(O(\log n)\) Tests whether \(k\) is not in the set.
--
-- @since 1.1.0.0
{-# INLINE notMember #-}
notMember :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool
notMember :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
notMember IntSet (PrimState m)
dis Int
k = Bool -> Bool
not (Bool -> Bool) -> m Bool -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet (PrimState m) -> Int -> m Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
member IntSet (PrimState m)
dis Int
k

-- | \(O(\log n)\) Looks up the smallest key \(k\) such that \(k \ge k_0\).
--
-- @since 1.1.0.0
{-# INLINE lookupGE #-}
lookupGE :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupGE :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupGE IntSet {Int
MVector (PrimState m) Int
Vector (MVector (PrimState m) Int)
capacityIS :: forall s. IntSet s -> Int
sizeIS :: forall s. IntSet s -> MVector s Int
vecIS :: forall s. IntSet s -> Vector (MVector s Int)
capacityIS :: Int
sizeIS :: MVector (PrimState m) Int
vecIS :: Vector (MVector (PrimState m) Int)
..} Int
i0
  | Int
i0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
capacityIS = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
  | Bool
otherwise = Int -> Int -> m (Maybe Int)
inner Int
0 (Int -> m (Maybe Int)) -> Int -> m (Maybe Int)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
i0 -- REMARK: it's very important to keep @i@ non-negative.
  where
    inner :: Int -> Int -> m (Maybe Int)
inner Int
h Int
i
      | Int
h Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector (MVector (PrimState m) Int) -> Int
forall a. Vector a -> Int
V.length Vector (MVector (PrimState m) Int)
vecIS = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
      -- ?
      | Int
q Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== MVector (PrimState m) Int -> Int
forall a s. Unbox a => MVector s a -> Int
VUM.length (Vector (MVector (PrimState m) Int)
-> Int -> MVector (PrimState m) Int
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex Vector (MVector (PrimState m) Int)
vecIS Int
h) = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
      | Bool
otherwise = do
          Int
d <- (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
r) (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead (Vector (MVector (PrimState m) Int)
-> Int -> MVector (PrimState m) Int
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex Vector (MVector (PrimState m) Int)
vecIS Int
h) Int
q
          if Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
            then Int -> Int -> m (Maybe Int)
inner (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            else
              Int -> Maybe Int
forall a. a -> Maybe a
Just
                (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> MVector (PrimState m) Int -> m Int)
-> Int -> Vector (MVector (PrimState m) Int) -> m Int
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Vector b -> m a
V.foldM'
                  ( \ !Int
acc MVector (PrimState m) Int
vec -> do
                      !Int
dx <- Int -> Int
lsbOf (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
vec Int
acc
                      Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dx
                  )
                  (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
lsbOf Int
d)
                  (Vector (MVector (PrimState m) Int)
-> Vector Int -> Vector (MVector (PrimState m) Int)
forall a. Vector a -> Vector Int -> Vector a
V.unsafeBackpermute Vector (MVector (PrimState m) Int)
vecIS (Int -> Int -> Int -> Vector Int
forall a. Num a => a -> a -> Int -> Vector a
V.enumFromStepN (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (-Int
1) Int
h))
      where
        (!Int
q, !Int
r) = Int
i Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
wordSize

-- | \(O(\log n)\) Looks up the smallest key \(k\) such that \(k \gt k_0\).
--
-- @since 1.1.0.0
{-# INLINE lookupGT #-}
lookupGT :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupGT :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupGT IntSet (PrimState m)
is Int
k = IntSet (PrimState m) -> Int -> m (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupGE IntSet (PrimState m)
is (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | \(O(\log n)\) Looks up the largest key \(k\) such that \(k \le k_0\).
--
-- @since 1.1.0.0
{-# INLINE lookupLE #-}
lookupLE :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupLE :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupLE IntSet {Int
MVector (PrimState m) Int
Vector (MVector (PrimState m) Int)
capacityIS :: forall s. IntSet s -> Int
sizeIS :: forall s. IntSet s -> MVector s Int
vecIS :: forall s. IntSet s -> Vector (MVector s Int)
capacityIS :: Int
sizeIS :: MVector (PrimState m) Int
vecIS :: Vector (MVector (PrimState m) Int)
..} Int
i0
  | Int
i0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= -Int
1 = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
  | Bool
otherwise = Int -> Int -> m (Maybe Int)
inner Int
0 (Int -> m (Maybe Int)) -> Int -> m (Maybe Int)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
capacityIS Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
i0
  where
    inner :: Int -> Int -> m (Maybe Int)
inner Int
h Int
i
      | Int
h Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector (MVector (PrimState m) Int) -> Int
forall a. Vector a -> Int
V.length Vector (MVector (PrimState m) Int)
vecIS = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1 = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
      | Bool
otherwise = do
          Int
d <- (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.<<. (Int
63 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r)) (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead (Vector (MVector (PrimState m) Int)
-> Int -> MVector (PrimState m) Int
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex Vector (MVector (PrimState m) Int)
vecIS Int
h) Int
q
          if Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
            then Int -> Int -> m (Maybe Int)
inner (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            else do
              Int -> Maybe Int
forall a. a -> Maybe a
Just
                (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> MVector (PrimState m) Int -> m Int)
-> Int -> Vector (MVector (PrimState m) Int) -> m Int
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Vector b -> m a
V.foldM'
                  ( \ !Int
acc MVector (PrimState m) Int
vec -> do
                      !Int
dx <- Int -> Int
msbOf (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
vec Int
acc
                      Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dx
                  )
                  (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Int
d)
                  (Vector (MVector (PrimState m) Int)
-> Vector Int -> Vector (MVector (PrimState m) Int)
forall a. Vector a -> Vector Int -> Vector a
V.unsafeBackpermute Vector (MVector (PrimState m) Int)
vecIS (Int -> Int -> Int -> Vector Int
forall a. Num a => a -> a -> Int -> Vector a
V.enumFromStepN (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (-Int
1) Int
h))
      where
        (!Int
q, !Int
r) = Int
i Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
wordSize

-- | \(O(\log n)\) Looks up the largest key \(k\) such that \(k \lt k_0\).
--
-- @since 1.1.0.0
{-# INLINE lookupLT #-}
lookupLT :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupLT :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupLT IntSet (PrimState m)
is Int
k = IntSet (PrimState m) -> Int -> m (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupLE IntSet (PrimState m)
is (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | \(O(\log n)\) Looks up the minimum key.
--
-- @since 1.1.0.0
{-# INLINE lookupMin #-}
lookupMin :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)
lookupMin :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> m (Maybe Int)
lookupMin IntSet (PrimState m)
is = IntSet (PrimState m) -> Int -> m (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupGE IntSet (PrimState m)
is Int
0

-- | \(O(\log n)\) Looks up the maximum key.
--
-- @since 1.1.0.0
{-# INLINE lookupMax #-}
lookupMax :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)
lookupMax :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> m (Maybe Int)
lookupMax IntSet (PrimState m)
is = IntSet (PrimState m) -> Int -> m (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupLE IntSet (PrimState m)
is (IntSet (PrimState m) -> Int
forall s. IntSet s -> Int
capacityIS IntSet (PrimState m)
is Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | \(O(\log n)\) Inserts a key \(k\) into the set. If an entry with the same key already exists,
-- it is overwritten.
--
-- @since 1.1.0.0
{-# INLINE insert #-}
insert :: (HasCallStack, PrimMonad m) => IntSet (PrimState m) -> Int -> m ()
insert :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
IntSet (PrimState m) -> Int -> m ()
insert is :: IntSet (PrimState m)
is@IntSet {Int
MVector (PrimState m) Int
Vector (MVector (PrimState m) Int)
capacityIS :: forall s. IntSet s -> Int
sizeIS :: forall s. IntSet s -> MVector s Int
vecIS :: forall s. IntSet s -> Vector (MVector s Int)
capacityIS :: Int
sizeIS :: MVector (PrimState m) Int
vecIS :: Vector (MVector (PrimState m) Int)
..} Int
k = do
  Bool
b <- IntSet (PrimState m) -> Int -> m Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
member IntSet (PrimState m)
is Int
k
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
b (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState m) Int -> (Int -> Int) -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
VUM.unsafeModify MVector (PrimState m) Int
sizeIS (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0
    (Int -> MVector (PrimState m) Int -> m Int)
-> Int -> Vector (MVector (PrimState m) Int) -> m ()
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Vector b -> m ()
V.foldM'_
      ( \Int
i MVector (PrimState m) Int
vec -> do
          let (!Int
q, !Int
r) = Int
i Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
wordSize
          MVector (PrimState m) Int -> (Int -> Int) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.unsafeModify MVector (PrimState m) Int
vec (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`setBit` Int
r) Int
q
          Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
q
      )
      Int
k
      Vector (MVector (PrimState m) Int)
vecIS
  where
    !()
_ = HasCallStack => String -> Int -> Int -> ()
String -> Int -> Int -> ()
ACIA.checkIndex String
"AtCoder.Extra.IntSet.insert" Int
k Int
capacityIS

-- | \(O(\log n)\) Deletes a key \(k\) from the set. Does nothing if no such key exists. Returns
-- whether the key existed.
--
-- @since 1.1.0.0
{-# INLINE delete #-}
delete :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool
delete :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
delete is :: IntSet (PrimState m)
is@IntSet {Int
MVector (PrimState m) Int
Vector (MVector (PrimState m) Int)
capacityIS :: forall s. IntSet s -> Int
sizeIS :: forall s. IntSet s -> MVector s Int
vecIS :: forall s. IntSet s -> Vector (MVector s Int)
capacityIS :: Int
sizeIS :: MVector (PrimState m) Int
vecIS :: Vector (MVector (PrimState m) Int)
..} Int
k = do
  Bool
b_ <- IntSet (PrimState m) -> Int -> m Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
member IntSet (PrimState m)
is Int
k
  if Bool
b_
    then do
      MVector (PrimState m) Int -> (Int -> Int) -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
VUM.unsafeModify MVector (PrimState m) Int
sizeIS (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) Int
0
      ((Bool, Int) -> MVector (PrimState m) Int -> m (Bool, Int))
-> (Bool, Int) -> Vector (MVector (PrimState m) Int) -> m ()
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Vector b -> m ()
V.foldM'_
        ( \(!Bool
b, !Int
i) MVector (PrimState m) Int
vec -> do
            let (!Int
q, !Int
r) = Int
i Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
wordSize
            -- TODO: early return is possible
            Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
b (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
              MVector (PrimState m) Int -> (Int -> Int) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.unsafeModify MVector (PrimState m) Int
vec (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`clearBit` Int
r) Int
q
            -- `b` remembers if any other bit was on
            Bool
b' <- (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (Int -> Bool) -> m Int -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
vec Int
q
            (Bool, Int) -> m (Bool, Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
b', Int
q)
        )
        (Bool
False, Int
k)
        Vector (MVector (PrimState m) Int)
vecIS
      Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
    else Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

-- | \(O(\log n)\) Deletes a key \(k\) from the set. Does nothing if no such key exists.
--
-- @since 1.1.0.0
{-# INLINE delete_ #-}
delete_ :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m ()
delete_ :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m ()
delete_ IntSet (PrimState m)
is Int
k = m Bool -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Bool -> m ()) -> m Bool -> m ()
forall a b. (a -> b) -> a -> b
$ IntSet (PrimState m) -> Int -> m Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
delete IntSet (PrimState m)
is Int
k

-- | \(O(\log n)\) Deletes the minimum key from the set. Returns `Nothing` if the set is empty.
--
-- @since 1.1.0.0
{-# INLINE deleteMin #-}
deleteMin :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)
deleteMin :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> m (Maybe Int)
deleteMin IntSet (PrimState m)
is = do
  IntSet (PrimState m) -> m (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> m (Maybe Int)
lookupMin IntSet (PrimState m)
is
    m (Maybe Int) -> (Maybe Int -> m (Maybe Int)) -> m (Maybe Int)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Int -> m Int) -> Maybe Int -> m (Maybe Int)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Maybe a -> m (Maybe b)
mapM
      ( \Int
key -> do
          IntSet (PrimState m) -> Int -> m ()
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m ()
delete_ IntSet (PrimState m)
is Int
key
          Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
key
      )

-- | \(O(\log n)\) Deletes the maximum key from the set. Returns `Nothing` if the set is empty.
--
-- @since 1.1.0.0
{-# INLINE deleteMax #-}
deleteMax :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)
deleteMax :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> m (Maybe Int)
deleteMax IntSet (PrimState m)
is = do
  IntSet (PrimState m) -> m (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> m (Maybe Int)
lookupMax IntSet (PrimState m)
is
    m (Maybe Int) -> (Maybe Int -> m (Maybe Int)) -> m (Maybe Int)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Int -> m Int) -> Maybe Int -> m (Maybe Int)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Maybe a -> m (Maybe b)
mapM
      ( \Int
key -> do
          IntSet (PrimState m) -> Int -> m ()
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m ()
delete_ IntSet (PrimState m)
is Int
key
          Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
key
      )

-- | \(O(n \log n)\) Enumerates the keys in the map.
--
-- @since 1.1.0.0
{-# INLINE keys #-}
keys :: (PrimMonad m) => IntSet (PrimState m) -> m (VU.Vector Int)
keys :: forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> m (Vector Int)
keys is :: IntSet (PrimState m)
is@IntSet {MVector (PrimState m) Int
sizeIS :: forall s. IntSet s -> MVector s Int
sizeIS :: MVector (PrimState m) Int
sizeIS} = do
  Int
n <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
sizeIS Int
0
  Int -> (Int -> m (Int, Int)) -> Int -> m (Vector Int)
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Int -> (b -> m (a, b)) -> b -> m (Vector a)
VU.unfoldrExactNM
    Int
n
    ( \Int
i -> do
        Int
i' <- Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> m (Maybe Int) -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet (PrimState m) -> Int -> m (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
lookupGT IntSet (PrimState m)
is Int
i
        (Int, Int) -> m (Int, Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i', Int
i')
    )
    (-Int
1)