{-# LANGUAGE RecordWildCards #-}

-- | A fast, mutable multiset for `Int` keys backed by a @HashMap@.  Most operations are performed
-- in \(O(1)\) time, but in average.
--
-- ==== Invariant
-- The count for each key must be non-negative. An exception is thrown if this invariant is
-- violated.
--
-- ==== Capacity limitation
-- The maximum number of distinct keys that can be inserted is fixed at \(n\), even if some keys are
-- deleted. This is due to the limitation of the internal @HashMap@.
--
-- ==== __Example__
-- Create a `MultiSet` with capacity \(4\):
--
-- >>> import AtCoder.Extra.MultiSet qualified as MS
-- >>> ms <- MS.new 4
--
-- `inc` and `dec` are the primary API:
--
-- >>> MS.inc ms 10
-- >>> MS.inc ms 10
-- >>> MS.lookup ms 10
-- Just 2
--
-- >>> MS.dec ms 10
-- >>> MS.lookup ms 10
-- Just 1
--
-- Entries with zero count are considered to be non-existing:
--
-- >>> MS.dec ms 10
-- >>> MS.member ms 10
-- False
--
-- >>> MS.lookup ms 10
-- Nothing
--
-- >>> MS.size ms
-- 0
--
-- Creating a negative count results in an exception:
--
-- >>> MS.inc ms 11
-- >>> MS.sub ms 11 2
-- *** Exception: AtCoder.Extra.Multiset.sub: the count of `11` is becoming a negative value: `-1`
-- ...
--
-- Decrementing a non-existing key does nothing and does not throw an exception:
--
-- >>> MS.dec ms 12
--
-- Misc:
--
-- >>> MS.insert ms 12 112
-- >>> MS.assocs ms
-- [(11,1),(12,112)]
--
-- @since 1.1.0.0
module AtCoder.Extra.MultiSet
  ( -- * MultiSet
    MultiSet,

    -- * Construtors
    new,

    -- * Metadata
    capacity,
    size,

    -- * Lookups
    lookup,
    member,
    notMember,

    -- * Modifications
    inc,
    dec,
    add,
    sub,
    insert,
    delete,

    -- * Conversions

    -- ** Safe conversions
    keys,
    elems,
    assocs,

    -- ** Unsafe conversions
    unsafeKeys,
    unsafeElems,
    unsafeAssocs,
  )
where

import AtCoder.Extra.HashMap qualified as HM
import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Functor ((<&>))
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 (lookup)

-- | A fast, mutable multiset for `Int` keys backed by a @HashMap@.
--
-- @since 1.1.0.0
data MultiSet s = MultiSet
  { forall s. MultiSet s -> HashMap s Int
mapMS :: !(HM.HashMap s Int),
    forall s. MultiSet s -> MVector s Int
cntMS :: !(VUM.MVector s Int)
  }

-- | \(O(n)\) Creates a `MultiSet` with capacity \(n\).
--
-- @since 1.1.0.0
new :: (PrimMonad m) => Int -> m (MultiSet (PrimState m))
new :: forall (m :: * -> *).
PrimMonad m =>
Int -> m (MultiSet (PrimState m))
new Int
n = do
  HashMap (PrimState m) Int
mapMS <- Int -> m (HashMap (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (HashMap (PrimState m) a)
HM.new Int
n
  MVector (PrimState m) Int
cntMS <- 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
  MultiSet (PrimState m) -> m (MultiSet (PrimState m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MultiSet (PrimState m) -> m (MultiSet (PrimState m)))
-> MultiSet (PrimState m) -> m (MultiSet (PrimState m))
forall a b. (a -> b) -> a -> b
$ MultiSet {MVector (PrimState m) Int
HashMap (PrimState m) Int
mapMS :: HashMap (PrimState m) Int
cntMS :: MVector (PrimState m) Int
mapMS :: HashMap (PrimState m) Int
cntMS :: MVector (PrimState m) Int
..}

-- | \(O(1)\) Returns the maximum number of distinct keys that can be inserted into the internal
-- hash map.
--
-- @since 1.1.0.0
capacity :: MultiSet s -> Int
capacity :: forall s. MultiSet s -> Int
capacity = HashMap s Int -> Int
forall s a. HashMap s a -> Int
HM.capacity (HashMap s Int -> Int)
-> (MultiSet s -> HashMap s Int) -> MultiSet s -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MultiSet s -> HashMap s Int
forall s. MultiSet s -> HashMap s Int
mapMS

-- | \(O(1)\) Returns the number of distinct keys with positive counts.
--
-- @since 1.1.0.0
size :: (PrimMonad m) => MultiSet (PrimState m) -> m Int
size :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m Int
size MultiSet {MVector (PrimState m) Int
HashMap (PrimState m) Int
mapMS :: forall s. MultiSet s -> HashMap s Int
cntMS :: forall s. MultiSet s -> MVector s Int
mapMS :: HashMap (PrimState m) Int
cntMS :: MVector (PrimState m) Int
..} = do
  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
cntMS Int
0

-- | \(O(1)\) Looks up the count for a key.
--
-- @since 1.1.0.0
lookup :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m (Maybe Int)
lookup :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> Int -> m (Maybe Int)
lookup MultiSet {MVector (PrimState m) Int
HashMap (PrimState m) Int
mapMS :: forall s. MultiSet s -> HashMap s Int
cntMS :: forall s. MultiSet s -> MVector s Int
mapMS :: HashMap (PrimState m) Int
cntMS :: MVector (PrimState m) Int
..} Int
k = do
  HashMap (PrimState m) Int -> Int -> m (Maybe Int)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap (PrimState m) Int
mapMS Int
k m (Maybe Int) -> (Maybe Int -> Maybe Int) -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
    Just Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i
    Maybe Int
_ -> Maybe Int
forall a. Maybe a
Nothing

-- | \(O(1)\) Tests whether \(k\) is in the set.
--
-- @since 1.1.0.0
member :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m Bool
member :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> Int -> m Bool
member MultiSet {MVector (PrimState m) Int
HashMap (PrimState m) Int
mapMS :: forall s. MultiSet s -> HashMap s Int
cntMS :: forall s. MultiSet s -> MVector s Int
mapMS :: HashMap (PrimState m) Int
cntMS :: MVector (PrimState m) Int
..} Int
k = do
  HashMap (PrimState m) Int -> Int -> m (Maybe Int)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap (PrimState m) Int
mapMS Int
k m (Maybe Int) -> (Maybe Int -> Bool) -> m Bool
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
    Just Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
    Maybe Int
_ -> Bool
False

-- | \(O(1)\) Tests whether \(k\) is not in the set.
--
-- @since 1.1.0.0
notMember :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m Bool
notMember :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> Int -> m Bool
notMember MultiSet (PrimState m)
ms Int
k = Bool -> Bool
not (Bool -> Bool) -> m Bool -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MultiSet (PrimState m) -> Int -> m Bool
forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> Int -> m Bool
member MultiSet (PrimState m)
ms Int
k

-- | \(O(1)\) Increments the count of a key.
--
-- @since 1.1.0.0
inc :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m ()
inc :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> m ()
inc MultiSet (PrimState m)
ms Int
k = MultiSet (PrimState m) -> Int -> Int -> m ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> Int -> m ()
add MultiSet (PrimState m)
ms Int
k Int
1

-- | \(O(1)\) Decrements the count of a key.
--
-- @since 1.1.0.0
dec :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m ()
dec :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> m ()
dec MultiSet (PrimState m)
ms Int
k = MultiSet (PrimState m) -> Int -> Int -> m ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> Int -> m ()
sub MultiSet (PrimState m)
ms Int
k Int
1

-- | \(O(1)\) Increments the count of a key \(k\) by \(c\). If the key does not exist in the set,
-- the \((k, c)\) pair is inserted. If \(v\) is negative, it falls back to `sub`.
--
-- @since 1.1.0.0
add :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m ()
add :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> Int -> m ()
add ms :: MultiSet (PrimState m)
ms@MultiSet {MVector (PrimState m) Int
HashMap (PrimState m) Int
mapMS :: forall s. MultiSet s -> HashMap s Int
cntMS :: forall s. MultiSet s -> MVector s Int
mapMS :: HashMap (PrimState m) Int
cntMS :: MVector (PrimState m) Int
..} Int
k Int
v = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
v Int
0 of
  Ordering
LT -> MultiSet (PrimState m) -> Int -> Int -> m ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> Int -> m ()
sub MultiSet (PrimState m)
ms Int
k (-Int
v)
  Ordering
EQ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Ordering
GT -> do
    HashMap (PrimState m) Int -> Int -> m (Maybe Int)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap (PrimState m) Int
mapMS Int
k m (Maybe Int) -> (Maybe Int -> 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
>>= \case
      Just Int
n ->  do
        HashMap (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap (PrimState m) Int
mapMS Int
k (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
v
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0) (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
cntMS (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0
      Maybe Int
Nothing -> do
        HashMap (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap (PrimState m) Int
mapMS Int
k Int
v
        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
cntMS (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0

-- | \(O(1)\) Decrements the count of a key \(k\) by \(c\). If \(c\) is negative, it falls back to
-- `add`.
--
-- @since 1.1.0.0
sub :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m ()
sub :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> Int -> m ()
sub ms :: MultiSet (PrimState m)
ms@MultiSet {MVector (PrimState m) Int
HashMap (PrimState m) Int
mapMS :: forall s. MultiSet s -> HashMap s Int
cntMS :: forall s. MultiSet s -> MVector s Int
mapMS :: HashMap (PrimState m) Int
cntMS :: MVector (PrimState m) Int
..} Int
k Int
v = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
v Int
0 of
  Ordering
LT -> MultiSet (PrimState m) -> Int -> Int -> m ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> Int -> m ()
add MultiSet (PrimState m)
ms Int
k (-Int
v)
  Ordering
EQ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Ordering
GT -> do
    HashMap (PrimState m) Int -> Int -> m (Maybe Int)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap (PrimState m) Int
mapMS Int
k m (Maybe Int) -> (Maybe Int -> 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
>>= \case
      Just Int
0 -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure () -- ignored
      Just Int
n -> case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
n Int
v of
        Ordering
GT -> do
          HashMap (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap (PrimState m) Int
mapMS Int
k (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
v)
        Ordering
EQ -> do
          HashMap (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap (PrimState m) Int
mapMS Int
k Int
0
          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
cntMS (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) Int
0
        Ordering
LT -> [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> m ()) -> [Char] -> m ()
forall a b. (a -> b) -> a -> b
$ [Char]
"AtCoder.Extra.Multiset.sub: the count of `" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
k [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"` is becoming a negative value: `" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
v) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"`"
      Maybe Int
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | \(O(1)\) Inserts a key-count pair into the set. `MultiSet` is actually a count map.
--
-- @since 1.1.0.0
insert :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m ()
insert :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> Int -> m ()
insert MultiSet {MVector (PrimState m) Int
HashMap (PrimState m) Int
mapMS :: forall s. MultiSet s -> HashMap s Int
cntMS :: forall s. MultiSet s -> MVector s Int
mapMS :: HashMap (PrimState m) Int
cntMS :: MVector (PrimState m) Int
..} Int
k Int
v
  | Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> m ()) -> [Char] -> m ()
forall a b. (a -> b) -> a -> b
$ [Char]
"AtCoder.Extra.Multiset.insert: new count must be positive`" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
k [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"`: `" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
v [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"`"
  | Bool
otherwise = do
      HashMap (PrimState m) Int -> Int -> m (Maybe Int)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap (PrimState m) Int
mapMS Int
k m (Maybe Int) -> (Maybe Int -> 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
>>= \case
        Just Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> do
          HashMap (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap (PrimState m) Int
mapMS Int
k Int
v
        Maybe Int
_ -> do
          HashMap (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap (PrimState m) Int
mapMS Int
k Int
v
          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
cntMS (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0

-- | \(O(1)\) Deletes a key. Note that it does not undo its insertion and does not increase the
-- number of distinct keys that can be inserted into the internal hash map.
--
-- @since 1.1.0.0
delete :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m ()
delete :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MultiSet (PrimState m) -> Int -> m ()
delete MultiSet {MVector (PrimState m) Int
HashMap (PrimState m) Int
mapMS :: forall s. MultiSet s -> HashMap s Int
cntMS :: forall s. MultiSet s -> MVector s Int
mapMS :: HashMap (PrimState m) Int
cntMS :: MVector (PrimState m) Int
..} Int
k = do
  HashMap (PrimState m) Int -> Int -> m (Maybe Int)
forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
HM.lookup HashMap (PrimState m) Int
mapMS Int
k m (Maybe Int) -> (Maybe Int -> 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
>>= \case
    Just Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> do
      HashMap (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
HM.insert HashMap (PrimState m) Int
mapMS Int
k Int
0
      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
cntMS (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) Int
0
    Maybe Int
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | \(O(n)\) Enumerates the keys in the set.
--
-- @since 1.1.0.0
{-# INLINE keys #-}
keys :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int)
keys :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m (Vector Int)
keys MultiSet (PrimState m)
ms = Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.force (Vector Int -> Vector Int) -> m (Vector Int) -> m (Vector Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MultiSet (PrimState m) -> m (Vector Int)
forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m (Vector Int)
unsafeKeys MultiSet (PrimState m)
ms

-- | \(O(n)\) Enumerates the counts in the set.
--
-- @since 1.1.0.0
{-# INLINE elems #-}
elems :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int)
elems :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m (Vector Int)
elems MultiSet (PrimState m)
ms = Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.force (Vector Int -> Vector Int) -> m (Vector Int) -> m (Vector Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MultiSet (PrimState m) -> m (Vector Int)
forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m (Vector Int)
unsafeElems MultiSet (PrimState m)
ms

-- | \(O(n)\) Enumerates the key-count pairs in the set.
--
-- @since 1.1.0.0
{-# INLINE assocs #-}
assocs :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector (Int, Int))
assocs :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m (Vector (Int, Int))
assocs MultiSet (PrimState m)
ms = Vector (Int, Int) -> Vector (Int, Int)
forall a. Unbox a => Vector a -> Vector a
VU.force (Vector (Int, Int) -> Vector (Int, Int))
-> m (Vector (Int, Int)) -> m (Vector (Int, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MultiSet (PrimState m) -> m (Vector (Int, Int))
forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m (Vector (Int, Int))
unsafeAssocs MultiSet (PrimState m)
ms

-- | \(O(n)\) Enumerates the keys in the set.
--
-- @since 1.1.0.0
{-# INLINE unsafeKeys #-}
unsafeKeys :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int)
unsafeKeys :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m (Vector Int)
unsafeKeys = (((Int, Int) -> Maybe Int) -> Vector (Int, Int) -> Vector Int
forall a b.
(Unbox a, Unbox b) =>
(a -> Maybe b) -> Vector a -> Vector b
VU.mapMaybe (\(!Int
k, !Int
n) -> if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then Int -> Maybe Int
forall a. a -> Maybe a
Just Int
k else Maybe Int
forall a. Maybe a
Nothing) <$>) (m (Vector (Int, Int)) -> m (Vector Int))
-> (MultiSet (PrimState m) -> m (Vector (Int, Int)))
-> MultiSet (PrimState m)
-> m (Vector Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap (PrimState m) Int -> m (Vector (Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector (Int, a))
HM.unsafeAssocs (HashMap (PrimState m) Int -> m (Vector (Int, Int)))
-> (MultiSet (PrimState m) -> HashMap (PrimState m) Int)
-> MultiSet (PrimState m)
-> m (Vector (Int, Int))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MultiSet (PrimState m) -> HashMap (PrimState m) Int
forall s. MultiSet s -> HashMap s Int
mapMS

-- | \(O(n)\) Enumerates the counts in the set.
--
-- @since 1.1.0.0
{-# INLINE unsafeElems #-}
unsafeElems :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int)
unsafeElems :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m (Vector Int)
unsafeElems = ((Int -> Bool) -> Vector Int -> Vector Int
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
VU.filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) <$>) (m (Vector Int) -> m (Vector Int))
-> (MultiSet (PrimState m) -> m (Vector Int))
-> MultiSet (PrimState m)
-> m (Vector Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap (PrimState m) Int -> m (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector a)
HM.unsafeElems (HashMap (PrimState m) Int -> m (Vector Int))
-> (MultiSet (PrimState m) -> HashMap (PrimState m) Int)
-> MultiSet (PrimState m)
-> m (Vector Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MultiSet (PrimState m) -> HashMap (PrimState m) Int
forall s. MultiSet s -> HashMap s Int
mapMS

-- | \(O(n)\) Enumerates the key-count pairs in the set.
--
-- @since 1.1.0.0
{-# INLINE unsafeAssocs #-}
unsafeAssocs :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector (Int, Int))
unsafeAssocs :: forall (m :: * -> *).
PrimMonad m =>
MultiSet (PrimState m) -> m (Vector (Int, Int))
unsafeAssocs = (((Int, Int) -> Bool) -> Vector (Int, Int) -> Vector (Int, Int)
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
VU.filter (\(!Int
_, !Int
n) -> Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) <$>) (m (Vector (Int, Int)) -> m (Vector (Int, Int)))
-> (MultiSet (PrimState m) -> m (Vector (Int, Int)))
-> MultiSet (PrimState m)
-> m (Vector (Int, Int))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap (PrimState m) Int -> m (Vector (Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector (Int, a))
HM.unsafeAssocs (HashMap (PrimState m) Int -> m (Vector (Int, Int)))
-> (MultiSet (PrimState m) -> HashMap (PrimState m) Int)
-> MultiSet (PrimState m)
-> m (Vector (Int, Int))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MultiSet (PrimState m) -> HashMap (PrimState m) Int
forall s. MultiSet s -> HashMap s Int
mapMS