{-# LANGUAGE TypeFamilies #-}

module Data.Multiset where

import Data.Coerce
import qualified Data.List as L
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as M
import GHC.Exts

newtype Multiset a = Multiset {forall a. Multiset a -> Map a Int
getMultiset :: M.Map a Int}
    deriving (Multiset a -> Multiset a -> Bool
(Multiset a -> Multiset a -> Bool)
-> (Multiset a -> Multiset a -> Bool) -> Eq (Multiset a)
forall a. Eq a => Multiset a -> Multiset a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Multiset a -> Multiset a -> Bool
== :: Multiset a -> Multiset a -> Bool
$c/= :: forall a. Eq a => Multiset a -> Multiset a -> Bool
/= :: Multiset a -> Multiset a -> Bool
Eq)

instance (Show a, Ord a) => Show (Multiset a) where
    show :: Multiset a -> String
show = [a] -> String
forall a. Show a => a -> String
show ([a] -> String) -> (Multiset a -> [a]) -> Multiset a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Multiset a -> [a]
Multiset a -> [Item (Multiset a)]
forall l. IsList l => l -> [Item l]
toList

instance (Ord a) => IsList (Multiset a) where
    type Item (Multiset a) = a
    fromList :: [Item (Multiset a)] -> Multiset a
fromList = [a] -> Multiset a
[Item (Multiset a)] -> Multiset a
forall a. Ord a => [a] -> Multiset a
fromListMS
    toList :: Multiset a -> [Item (Multiset a)]
toList = Multiset a -> [a]
Multiset a -> [Item (Multiset a)]
forall a. Multiset a -> [a]
toListMS

emptyMS :: Multiset a
emptyMS :: forall a. Multiset a
emptyMS = Map a Int -> Multiset a
forall a b. Coercible a b => a -> b
coerce (forall k a. Map k a
M.empty @_ @Int)

singletonMS :: a -> Multiset a
singletonMS :: forall a. a -> Multiset a
singletonMS = (a -> Map a Int) -> a -> Multiset a
forall a b. Coercible a b => a -> b
coerce ((a -> Int -> Map a Int) -> Int -> a -> Map a Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall k a. k -> a -> Map k a
M.singleton @_ @Int) Int
1)

{- | /O(1)/

>>> replicateMS 3 1
[1,1,1]
>>> nullMS $ replicateMS 0 1
True
>>> nullMS $ replicateMS (-1) 1
True
-}
replicateMS :: Int -> a -> Multiset a
replicateMS :: forall a. Int -> a -> Multiset a
replicateMS = (Int -> a -> Map a Int) -> Int -> a -> Multiset a
forall a b. Coercible a b => a -> b
coerce (((Int -> Bool) -> Map a Int -> Map a Int
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Map a Int -> Map a Int) -> (a -> Map a Int) -> a -> Map a Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((a -> Map a Int) -> a -> Map a Int)
-> (Int -> a -> Map a Int) -> Int -> a -> Map a Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Int -> Map a Int) -> Int -> a -> Map a Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall k a. k -> a -> Map k a
M.singleton @_ @Int))

{- | /O(n log n)/

>>> fromListMS [0,1,2,1,0]
[0,0,1,1,2]
-}
fromListMS :: (Ord a) => [a] -> Multiset a
fromListMS :: forall a. Ord a => [a] -> Multiset a
fromListMS = (Multiset a -> a -> Multiset a) -> Multiset a -> [a] -> Multiset a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' ((a -> Multiset a -> Multiset a) -> Multiset a -> a -> Multiset a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> Multiset a -> Multiset a
forall a. Ord a => a -> Multiset a -> Multiset a
insertMS) Multiset a
forall a. Multiset a
emptyMS

{- | /O(n)/

>>> fromAscListMS [0,0,1,2]
[0,0,1,2]
-}
fromAscListMS :: (Eq a) => [a] -> Multiset a
fromAscListMS :: forall a. Eq a => [a] -> Multiset a
fromAscListMS =
    Map a Int -> Multiset a
forall a. Map a Int -> Multiset a
Multiset
        (Map a Int -> Multiset a)
-> ([a] -> Map a Int) -> [a] -> Multiset a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, Int)] -> Map a Int
forall k a. [(k, a)] -> Map k a
M.fromDistinctAscList
        ([(a, Int)] -> Map a Int)
-> ([a] -> [(a, Int)]) -> [a] -> Map a Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NonEmpty a -> (a, Int)) -> [NonEmpty a] -> [(a, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\NonEmpty a
g -> (NonEmpty a -> a
forall a. NonEmpty a -> a
NE.head NonEmpty a
g, NonEmpty a -> Int
forall a. NonEmpty a -> Int
NE.length NonEmpty a
g))
        ([NonEmpty a] -> [(a, Int)])
-> ([a] -> [NonEmpty a]) -> [a] -> [(a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [NonEmpty a]
forall (f :: * -> *) a. (Foldable f, Eq a) => f a -> [NonEmpty a]
NE.group

{- | /O(n)/

>>> fromDistinctAscListMS [0,1,2]
[0,1,2]
-}
fromDistinctAscListMS :: [a] -> Multiset a
fromDistinctAscListMS :: forall a. [a] -> Multiset a
fromDistinctAscListMS =
    ([a] -> Map a Int) -> [a] -> Multiset a
forall a b. Coercible a b => a -> b
coerce (forall k a. [(k, a)] -> Map k a
M.fromDistinctAscList @_ @Int ([(a, Int)] -> Map a Int)
-> ([a] -> [(a, Int)]) -> [a] -> Map a Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (a, Int)) -> [a] -> [(a, Int)]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Int -> (a, Int)) -> Int -> a -> (a, Int)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) Int
1))

{- | /O(n)/

>>> fromDescListMS [2,1,0,0]
[0,0,1,2]
-}
fromDescListMS :: (Eq a) => [a] -> Multiset a
fromDescListMS :: forall a. Eq a => [a] -> Multiset a
fromDescListMS =
    Map a Int -> Multiset a
forall a. Map a Int -> Multiset a
Multiset
        (Map a Int -> Multiset a)
-> ([a] -> Map a Int) -> [a] -> Multiset a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, Int)] -> Map a Int
forall k a. [(k, a)] -> Map k a
M.fromDistinctAscList
        ([(a, Int)] -> Map a Int)
-> ([a] -> [(a, Int)]) -> [a] -> Map a Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, Int)] -> [(a, Int)]
forall a. [a] -> [a]
reverse
        ([(a, Int)] -> [(a, Int)])
-> ([a] -> [(a, Int)]) -> [a] -> [(a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NonEmpty a -> (a, Int)) -> [NonEmpty a] -> [(a, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\NonEmpty a
g -> (NonEmpty a -> a
forall a. NonEmpty a -> a
NE.head NonEmpty a
g, NonEmpty a -> Int
forall a. NonEmpty a -> Int
NE.length NonEmpty a
g))
        ([NonEmpty a] -> [(a, Int)])
-> ([a] -> [NonEmpty a]) -> [a] -> [(a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [NonEmpty a]
forall (f :: * -> *) a. (Foldable f, Eq a) => f a -> [NonEmpty a]
NE.group

{- | /O(n)/

>>> fromDistinctDescListMS [2,1,0]
[0,1,2]
-}
fromDistinctDescListMS :: [a] -> Multiset a
fromDistinctDescListMS :: forall a. [a] -> Multiset a
fromDistinctDescListMS =
    ([a] -> Map a Int) -> [a] -> Multiset a
forall a b. Coercible a b => a -> b
coerce (forall k a. [(k, a)] -> Map k a
M.fromDistinctAscList @_ @Int ([(a, Int)] -> Map a Int)
-> ([a] -> [(a, Int)]) -> [a] -> Map a Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (a, Int)) -> [a] -> [(a, Int)]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Int -> (a, Int)) -> Int -> a -> (a, Int)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) Int
1) ([a] -> [(a, Int)]) -> ([a] -> [a]) -> [a] -> [(a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse)

{- | /O(n)/

>>> toListMS (fromListMS [0,1,0,2])
[0,0,1,2]
-}
toListMS :: Multiset a -> [a]
toListMS :: forall a. Multiset a -> [a]
toListMS = Multiset a -> [a]
forall a. Multiset a -> [a]
toAscListMS

{- | /O(n)/

>>> toAscListMS (fromListMS [0,1,0,2])
[0,0,1,2]
-}
toAscListMS :: Multiset a -> [a]
toAscListMS :: forall a. Multiset a -> [a]
toAscListMS =
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap @[] (\(a
k, Int
x) -> Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
x a
k)
        ([(a, Int)] -> [a])
-> (Multiset a -> [(a, Int)]) -> Multiset a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map a Int -> [(a, Int)]) -> Multiset a -> [(a, Int)]
forall a b. Coercible a b => a -> b
coerce (forall k a. Map k a -> [(k, a)]
M.toAscList @_ @Int)

{- | /O(n)/

>>> toDescListMS (fromListMS [0,1,0,2])
[2,1,0,0]
-}
toDescListMS :: Multiset a -> [a]
toDescListMS :: forall a. Multiset a -> [a]
toDescListMS =
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap @[] (\(a
k, Int
x) -> Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
x a
k)
        ([(a, Int)] -> [a])
-> (Multiset a -> [(a, Int)]) -> Multiset a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map a Int -> [(a, Int)]) -> Multiset a -> [(a, Int)]
forall a b. Coercible a b => a -> b
coerce (forall k a. Map k a -> [(k, a)]
M.toDescList @_ @Int)

-- | /O(log n)/
insertMS :: (Ord a) => a -> Multiset a -> Multiset a
insertMS :: forall a. Ord a => a -> Multiset a -> Multiset a
insertMS = (a -> Map a Int -> Map a Int) -> a -> Multiset a -> Multiset a
forall a b. Coercible a b => a -> b
coerce ((a -> Int -> Map a Int -> Map a Int)
-> Int -> a -> Map a Int -> Map a Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith @_ @Int Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)) Int
1)
{-# INLINE insertMS #-}

{- | /O(log n)/

>>> deleteMS 0 (fromList [0, 0, 1, 2])
[0,1,2]
>>> deleteMS 3 (fromList [0, 0, 1, 2])
[0,0,1,2]
-}
deleteMS :: (Ord a) => a -> Multiset a -> Multiset a
deleteMS :: forall a. Ord a => a -> Multiset a -> Multiset a
deleteMS = (a -> Map a Int -> Map a Int) -> a -> Multiset a -> Multiset a
forall a b. Coercible a b => a -> b
coerce (forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
M.update @_ @Int (\Int
y -> if Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 then Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$! Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 else Maybe Int
forall a. Maybe a
Nothing))

{- | /O(log n)/

>>> deleteAllMS 0 (fromList [0, 0, 1, 2])
[1,2]
>>> deleteAllMS 3 (fromList [0, 0, 1, 2])
[0,0,1,2]
-}
deleteAllMS :: (Ord a) => a -> Multiset a -> Multiset a
deleteAllMS :: forall a. Ord a => a -> Multiset a -> Multiset a
deleteAllMS = (a -> Map a Int -> Map a Int) -> a -> Multiset a -> Multiset a
forall a b. Coercible a b => a -> b
coerce (forall k a. Ord k => k -> Map k a -> Map k a
M.delete @_ @Int)

-- | /O(log n)/
memberMS :: (Ord a) => a -> Multiset a -> Bool
memberMS :: forall a. Ord a => a -> Multiset a -> Bool
memberMS = (a -> Map a Int -> Bool) -> a -> Multiset a -> Bool
forall a b. Coercible a b => a -> b
coerce (forall k a. Ord k => k -> Map k a -> Bool
M.member @_ @Int)

-- | /O(log n)/
notMemberMS :: (Ord a) => a -> Multiset a -> Bool
notMemberMS :: forall a. Ord a => a -> Multiset a -> Bool
notMemberMS = (a -> Map a Int -> Bool) -> a -> Multiset a -> Bool
forall a b. Coercible a b => a -> b
coerce (forall k a. Ord k => k -> Map k a -> Bool
M.notMember @_ @Int)

{- | /O(log n)/

>>> countMS 0 (fromList [0, 0, 1, 2])
2
>>> countMS 1 (fromList [0, 0, 1, 2])
1
>>> countMS (-1) (fromList [0, 0, 1, 2])
0
-}
countMS :: (Ord a) => a -> Multiset a -> Int
countMS :: forall a. Ord a => a -> Multiset a -> Int
countMS = (a -> Map a Int -> Int) -> a -> Multiset a -> Int
forall a b. Coercible a b => a -> b
coerce (forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault @_ @Int Int
0)

{- | /O(log n)/

>>> lookupLTMS 1 (fromList [0, 0, 1, 1, 2, 2])
Just 0
>>> lookupLTMS 0 (fromList [0, 0, 1, 1, 2, 2])
Nothing
-}
lookupLTMS :: (Ord a) => a -> Multiset a -> Maybe a
lookupLTMS :: forall a. Ord a => a -> Multiset a -> Maybe a
lookupLTMS a
x = (Map a Int -> Maybe a) -> Multiset a -> Maybe a
forall a b. Coercible a b => a -> b
coerce (((a, Int) -> a) -> Maybe (a, Int) -> Maybe a
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Int) -> a
forall a b. (a, b) -> a
fst (Maybe (a, Int) -> Maybe a)
-> (Map a Int -> Maybe (a, Int)) -> Map a Int -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. Ord k => k -> Map k v -> Maybe (k, v)
M.lookupLT @_ @Int a
x)

{- | /O(log n)/

>>> lookupGTMS 1 (fromList [0, 0, 1, 1, 2, 2])
Just 2
>>> lookupGTMS 2 (fromList [0, 0, 1, 1, 2, 2])
Nothing
-}
lookupGTMS :: (Ord a) => a -> Multiset a -> Maybe a
lookupGTMS :: forall a. Ord a => a -> Multiset a -> Maybe a
lookupGTMS a
x = (Map a Int -> Maybe a) -> Multiset a -> Maybe a
forall a b. Coercible a b => a -> b
coerce (((a, Int) -> a) -> Maybe (a, Int) -> Maybe a
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Int) -> a
forall a b. (a, b) -> a
fst (Maybe (a, Int) -> Maybe a)
-> (Map a Int -> Maybe (a, Int)) -> Map a Int -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. Ord k => k -> Map k v -> Maybe (k, v)
M.lookupGT @_ @Int a
x)

{- | /O(log n)/

>>> lookupLEMS 1 (fromList [0, 0, 1, 1, 2, 2])
Just 1
>>> lookupLEMS 0 (fromList [0, 0, 1, 1, 2, 2])
Just 0
>>> lookupLEMS (-1) (fromList [0, 0, 1, 1, 2, 2])
Nothing
-}
lookupLEMS :: (Ord a) => a -> Multiset a -> Maybe a
lookupLEMS :: forall a. Ord a => a -> Multiset a -> Maybe a
lookupLEMS a
x = (Map a Int -> Maybe a) -> Multiset a -> Maybe a
forall a b. Coercible a b => a -> b
coerce (((a, Int) -> a) -> Maybe (a, Int) -> Maybe a
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Int) -> a
forall a b. (a, b) -> a
fst (Maybe (a, Int) -> Maybe a)
-> (Map a Int -> Maybe (a, Int)) -> Map a Int -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. Ord k => k -> Map k v -> Maybe (k, v)
M.lookupLE @_ @Int a
x)

{- | /O(log n)/

>>> lookupGEMS 1 (fromList [0, 0, 1, 1, 2, 2])
Just 1
>>> lookupGEMS 2 (fromList [0, 0, 1, 1, 2, 2])
Just 2
>>> lookupGEMS 3 (fromList [0, 0, 1, 1, 2, 2])
Nothing
-}
lookupGEMS :: (Ord a) => a -> Multiset a -> Maybe a
lookupGEMS :: forall a. Ord a => a -> Multiset a -> Maybe a
lookupGEMS a
x = (Map a Int -> Maybe a) -> Multiset a -> Maybe a
forall a b. Coercible a b => a -> b
coerce (((a, Int) -> a) -> Maybe (a, Int) -> Maybe a
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Int) -> a
forall a b. (a, b) -> a
fst (Maybe (a, Int) -> Maybe a)
-> (Map a Int -> Maybe (a, Int)) -> Map a Int -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. Ord k => k -> Map k v -> Maybe (k, v)
M.lookupGE @_ @Int a
x)

-- | /O(1)/
nullMS :: Multiset a -> Bool
nullMS :: forall a. Multiset a -> Bool
nullMS = (Map a Int -> Bool) -> Multiset a -> Bool
forall a b. Coercible a b => a -> b
coerce (forall k a. Map k a -> Bool
M.null @_ @Int)

{- | /O(n)/

>>> sizeMS (fromList [0, 0, 1, 2])
4
>>> sizeMS emptyMS
0
-}
sizeMS :: Multiset a -> Int
sizeMS :: forall a. Multiset a -> Int
sizeMS = (Map a Int -> Int) -> Multiset a -> Int
forall a b. Coercible a b => a -> b
coerce (forall a b k. (a -> b -> a) -> a -> Map k b -> a
M.foldl' @Int Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0)

{- | /O(1)/
>>> distinctSizeMS (fromList [1,1,1,2,2,3])
3
-}
distinctSizeMS :: Multiset a -> Int
distinctSizeMS :: forall a. Multiset a -> Int
distinctSizeMS = (Map a Int -> Int) -> Multiset a -> Int
forall a b. Coercible a b => a -> b
coerce (forall k a. Map k a -> Int
M.size @_ @Int)

{- | /O(log n)/

>>> findMinMS (fromList [0, 0, 1, 2])
0
>>> findMinMS emptyMS
*** Exception: Map.findMin: empty map has no minimal element
-}
findMinMS :: Multiset a -> a
findMinMS :: forall a. Multiset a -> a
findMinMS = (Map a Int -> a) -> Multiset a -> a
forall a b. Coercible a b => a -> b
coerce ((a, Int) -> a
forall a b. (a, b) -> a
fst ((a, Int) -> a) -> (Map a Int -> (a, Int)) -> Map a Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> (k, a)
M.findMin @_ @Int)

{- | /O(log n)/

>>> findMaxMS (fromList [0, 0, 1, 2])
2

>>> findMaxMS emptyMS
*** Exception: Map.findMax: empty map has no maximal element
-}
findMaxMS :: Multiset a -> a
findMaxMS :: forall a. Multiset a -> a
findMaxMS = (Map a Int -> a) -> Multiset a -> a
forall a b. Coercible a b => a -> b
coerce ((a, Int) -> a
forall a b. (a, b) -> a
fst ((a, Int) -> a) -> (Map a Int -> (a, Int)) -> Map a Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> (k, a)
M.findMax @_ @Int)

{- | /O(log n)/

>>> deleteMinMS (fromList [0, 0, 1, 2])
[0,1,2]
>>> deleteMinMS emptyMS
[]
-}
deleteMinMS :: Multiset a -> Multiset a
deleteMinMS :: forall a. Multiset a -> Multiset a
deleteMinMS =
    (Map a Int -> Map a Int) -> Multiset a -> Multiset a
forall a b. Coercible a b => a -> b
coerce (forall a k. (a -> Maybe a) -> Map k a -> Map k a
M.updateMin @Int (\Int
x -> if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 then Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$! Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 else Maybe Int
forall a. Maybe a
Nothing))

{- | /O(log n)/

>>> deleteMaxMS (fromList [0, 1, 2, 2])
[0,1,2]
>>> deleteMaxMS emptyMS
[]
-}
deleteMaxMS :: Multiset a -> Multiset a
deleteMaxMS :: forall a. Multiset a -> Multiset a
deleteMaxMS =
    (Map a Int -> Map a Int) -> Multiset a -> Multiset a
forall a b. Coercible a b => a -> b
coerce (forall a k. (a -> Maybe a) -> Map k a -> Map k a
M.updateMax @Int (\Int
x -> if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 then Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$! Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 else Maybe Int
forall a. Maybe a
Nothing))

{- | /O(log n)/

>>> deleteFindMinMS (fromList [0, 0, 1, 2])
(0,[0,1,2])
>>> deleteFindMinMS emptyMS
*** Exception: Map.deleteFindMin: can not return the minimal element of an empty map
-}
deleteFindMinMS :: forall a. (Ord a) => Multiset a -> (a, Multiset a)
deleteFindMinMS :: forall a. Ord a => Multiset a -> (a, Multiset a)
deleteFindMinMS = (Map a Int -> (a, Map a Int)) -> Multiset a -> (a, Multiset a)
forall a b. Coercible a b => a -> b
coerce (((a, Int), Map a Int) -> (a, Map a Int)
found (((a, Int), Map a Int) -> (a, Map a Int))
-> (Map a Int -> ((a, Int), Map a Int))
-> Map a Int
-> (a, Map a Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a Int -> ((a, Int), Map a Int)
forall k a. Map k a -> ((k, a), Map k a)
M.deleteFindMin)
    where
        found :: ((a, Int), M.Map a Int) -> (a, M.Map a Int)
        found :: ((a, Int), Map a Int) -> (a, Map a Int)
found ((a
k, Int
x), Map a Int
m)
            | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 = case a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
k (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Map a Int
m of
                Map a Int
m' -> (a
k, Map a Int
m')
            | Bool
otherwise = (a
k, Map a Int
m)

{- | /O(log n)/

>>> deleteFindMaxMS (fromList [0, 1, 2, 2])
(2,[0,1,2])

>>> deleteFindMaxMS emptyMS
*** Exception: Map.deleteFindMax: can not return the maximal element of an empty map
-}
deleteFindMaxMS :: forall a. (Ord a) => Multiset a -> (a, Multiset a)
deleteFindMaxMS :: forall a. Ord a => Multiset a -> (a, Multiset a)
deleteFindMaxMS = (Map a Int -> (a, Map a Int)) -> Multiset a -> (a, Multiset a)
forall a b. Coercible a b => a -> b
coerce (((a, Int), Map a Int) -> (a, Map a Int)
found (((a, Int), Map a Int) -> (a, Map a Int))
-> (Map a Int -> ((a, Int), Map a Int))
-> Map a Int
-> (a, Map a Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a Int -> ((a, Int), Map a Int)
forall k a. Map k a -> ((k, a), Map k a)
M.deleteFindMax)
    where
        found :: ((a, Int), M.Map a Int) -> (a, M.Map a Int)
        found :: ((a, Int), Map a Int) -> (a, Map a Int)
found ((a
k, Int
x), Map a Int
m)
            | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 = case a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
k (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Map a Int
m of
                Map a Int
m' -> (a
k, Map a Int
m')
            | Bool
otherwise = (a
k, Map a Int
m)

{- | /O(log n)/

>>> minViewMS (fromList [0, 0, 1, 2])
Just (0,[0,1,2])

>>> minViewMS emptyMS
Nothing
-}
minViewMS :: forall a. (Ord a) => Multiset a -> Maybe (a, Multiset a)
minViewMS :: forall a. Ord a => Multiset a -> Maybe (a, Multiset a)
minViewMS = (Map a Int -> Maybe (a, Multiset a))
-> Multiset a -> Maybe (a, Multiset a)
forall a b. Coercible a b => a -> b
coerce (Maybe (a, Multiset a)
-> (((a, Int), Map a Int) -> Maybe (a, Multiset a))
-> Maybe ((a, Int), Map a Int)
-> Maybe (a, Multiset a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Maybe (a, Multiset a)
forall a. Maybe a
Nothing ((a, Int), Map a Int) -> Maybe (a, Multiset a)
just (Maybe ((a, Int), Map a Int) -> Maybe (a, Multiset a))
-> (Map a Int -> Maybe ((a, Int), Map a Int))
-> Map a Int
-> Maybe (a, Multiset a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a Int -> Maybe ((a, Int), Map a Int)
forall k a. Map k a -> Maybe ((k, a), Map k a)
M.minViewWithKey)
    where
        just :: ((a, Int), M.Map a Int) -> Maybe (a, Multiset a)
        just :: ((a, Int), Map a Int) -> Maybe (a, Multiset a)
just ((a
k, Int
x), Map a Int
m)
            | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 = case a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
k (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Map a Int
m of
                Map a Int
m' -> Maybe (a, Map a Int) -> Maybe (a, Multiset a)
forall a b. Coercible a b => a -> b
coerce ((a, Map a Int) -> Maybe (a, Map a Int)
forall a. a -> Maybe a
Just (a
k, Map a Int
m'))
            | Bool
otherwise = Maybe (a, Map a Int) -> Maybe (a, Multiset a)
forall a b. Coercible a b => a -> b
coerce ((a, Map a Int) -> Maybe (a, Map a Int)
forall a. a -> Maybe a
Just (a
k, Map a Int
m))

{- | /O(log n)/

>>> maxViewMS (fromList [0, 1, 2, 2])
Just (2,[0,1,2])

>>> maxViewMS emptyMS
Nothing
-}
maxViewMS :: forall a. (Ord a) => Multiset a -> Maybe (a, Multiset a)
maxViewMS :: forall a. Ord a => Multiset a -> Maybe (a, Multiset a)
maxViewMS = (Map a Int -> Maybe (a, Multiset a))
-> Multiset a -> Maybe (a, Multiset a)
forall a b. Coercible a b => a -> b
coerce (Maybe (a, Multiset a)
-> (((a, Int), Map a Int) -> Maybe (a, Multiset a))
-> Maybe ((a, Int), Map a Int)
-> Maybe (a, Multiset a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Maybe (a, Multiset a)
forall a. Maybe a
Nothing ((a, Int), Map a Int) -> Maybe (a, Multiset a)
just (Maybe ((a, Int), Map a Int) -> Maybe (a, Multiset a))
-> (Map a Int -> Maybe ((a, Int), Map a Int))
-> Map a Int
-> Maybe (a, Multiset a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a Int -> Maybe ((a, Int), Map a Int)
forall k a. Map k a -> Maybe ((k, a), Map k a)
M.maxViewWithKey)
    where
        just :: ((a, Int), M.Map a Int) -> Maybe (a, Multiset a)
        just :: ((a, Int), Map a Int) -> Maybe (a, Multiset a)
just ((a
k, Int
x), Map a Int
m)
            | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 = case a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
k (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Map a Int
m of
                Map a Int
m' -> Maybe (a, Map a Int) -> Maybe (a, Multiset a)
forall a b. Coercible a b => a -> b
coerce ((a, Map a Int) -> Maybe (a, Map a Int)
forall a. a -> Maybe a
Just (a
k, Map a Int
m'))
            | Bool
otherwise = Maybe (a, Map a Int) -> Maybe (a, Multiset a)
forall a b. Coercible a b => a -> b
coerce ((a, Map a Int) -> Maybe (a, Map a Int)
forall a. a -> Maybe a
Just (a
k, Map a Int
m))

{- | /O(log n)/

>>> splitMS 1 (fromList [0, 0, 1, 2])
([0,0],[2])
>>> splitMS 0 (fromList [0, 0, 1, 2])
([],[1,2])
>>> splitMS (-1) (fromList [0, 0, 1, 2])
([],[0,0,1,2])
-}
splitMS :: (Ord a) => a -> Multiset a -> (Multiset a, Multiset a)
splitMS :: forall a. Ord a => a -> Multiset a -> (Multiset a, Multiset a)
splitMS = (a -> Map a Int -> (Map a Int, Map a Int))
-> a -> Multiset a -> (Multiset a, Multiset a)
forall a b. Coercible a b => a -> b
coerce (forall k a. Ord k => k -> Map k a -> (Map k a, Map k a)
M.split @_ @Int)

{- | /O(log n)/

>>> splitLookupMS 1 (fromList [0, 0, 1, 2])
([0,0],[1],[2])
>>> splitLookupMS 0 (fromList [0, 0, 1, 2])
([],[0,0],[1,2])
>>> splitLookupMS (-1) (fromList [0, 0, 1, 2])
([],[],[0,0,1,2])
-}
splitLookupMS :: (Ord a) => a -> Multiset a -> (Multiset a, [a], Multiset a)
splitLookupMS :: forall a. Ord a => a -> Multiset a -> (Multiset a, [a], Multiset a)
splitLookupMS a
k Multiset a
h = case (a -> Map a Int -> (Map a Int, Maybe Int, Map a Int))
-> a -> Multiset a -> (Multiset a, Maybe Int, Multiset a)
forall a b. Coercible a b => a -> b
coerce (forall k a. Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
M.splitLookup @_ @Int) a
k Multiset a
h of
    (Multiset a
l, Just Int
cnt, Multiset a
r) -> (Multiset a
l, Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
cnt a
k, Multiset a
r)
    (Multiset a
l, Maybe Int
Nothing, Multiset a
r) -> (Multiset a
l, [], Multiset a
r)