{-# LANGUAGE TypeFamilies #-}

module Data.IntHeap where

import Data.Coerce
import qualified Data.IntMap.Strict as IM
import qualified Data.List as L
import qualified Data.List.NonEmpty as NE
import GHC.Exts

newtype IntHeap = IntHeap {IntHeap -> IntMap Key
getIntHeap :: IM.IntMap Int}
  deriving (IntHeap -> IntHeap -> Bool
(IntHeap -> IntHeap -> Bool)
-> (IntHeap -> IntHeap -> Bool) -> Eq IntHeap
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: IntHeap -> IntHeap -> Bool
== :: IntHeap -> IntHeap -> Bool
$c/= :: IntHeap -> IntHeap -> Bool
/= :: IntHeap -> IntHeap -> Bool
Eq)

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

instance IsList IntHeap where
  type Item IntHeap = Int
  fromList :: [Item IntHeap] -> IntHeap
fromList = [Key] -> IntHeap
[Item IntHeap] -> IntHeap
fromListIH
  toList :: IntHeap -> [Item IntHeap]
toList = IntHeap -> [Key]
IntHeap -> [Item IntHeap]
toListIH

emptyIH :: IntHeap
emptyIH :: IntHeap
emptyIH = IntMap Key -> IntHeap
forall a b. Coercible a b => a -> b
coerce (forall a. IntMap a
IM.empty @Int)

singletonIH :: Int -> IntHeap
singletonIH :: Key -> IntHeap
singletonIH = (Key -> IntMap Key) -> Key -> IntHeap
forall a b. Coercible a b => a -> b
coerce ((Key -> Key -> IntMap Key) -> Key -> Key -> IntMap Key
forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall a. Key -> a -> IntMap a
IM.singleton @Int) Key
1)

{- | /O(1)/

>>> replicateIH 3 1
[1,1,1]
>>> nullIH $ replicateIH 0 1
True
>>> nullIH $ replicateIH (-1) 1
True
-}
replicateIH :: Int -> Int -> IntHeap
replicateIH :: Key -> Key -> IntHeap
replicateIH = (Key -> Key -> IntMap Key) -> Key -> Key -> IntHeap
forall a b. Coercible a b => a -> b
coerce (((Key -> Bool) -> IntMap Key -> IntMap Key
forall a. (a -> Bool) -> IntMap a -> IntMap a
IM.filter (Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
0) .) ((Key -> IntMap Key) -> Key -> IntMap Key)
-> (Key -> Key -> IntMap Key) -> Key -> Key -> IntMap Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Key -> Key -> IntMap Key) -> Key -> Key -> IntMap Key
forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall a. Key -> a -> IntMap a
IM.singleton @Int))

{- | /O(n min(n,W))/

>>> fromListIH [0,1,2,1,0]
[0,0,1,1,2]
-}
fromListIH :: [Int] -> IntHeap
fromListIH :: [Key] -> IntHeap
fromListIH = (IntHeap -> Key -> IntHeap) -> IntHeap -> [Key] -> IntHeap
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' ((Key -> IntHeap -> IntHeap) -> IntHeap -> Key -> IntHeap
forall a b c. (a -> b -> c) -> b -> a -> c
flip Key -> IntHeap -> IntHeap
insertIH) IntHeap
emptyIH

{- | /O(n)/

>>> fromAscListIH [0,0,1,2]
[0,0,1,2]
-}
fromAscListIH :: [Int] -> IntHeap
fromAscListIH :: [Key] -> IntHeap
fromAscListIH =
  IntMap Key -> IntHeap
IntHeap
    (IntMap Key -> IntHeap)
-> ([Key] -> IntMap Key) -> [Key] -> IntHeap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Key, Key)] -> IntMap Key
forall a. [(Key, a)] -> IntMap a
IM.fromDistinctAscList
    ([(Key, Key)] -> IntMap Key)
-> ([Key] -> [(Key, Key)]) -> [Key] -> IntMap Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NonEmpty Key -> (Key, Key)) -> [NonEmpty Key] -> [(Key, Key)]
forall a b. (a -> b) -> [a] -> [b]
map (\NonEmpty Key
g -> (NonEmpty Key -> Key
forall a. NonEmpty a -> a
NE.head NonEmpty Key
g, NonEmpty Key -> Key
forall a. NonEmpty a -> Key
NE.length NonEmpty Key
g))
    ([NonEmpty Key] -> [(Key, Key)])
-> ([Key] -> [NonEmpty Key]) -> [Key] -> [(Key, Key)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Key] -> [NonEmpty Key]
forall (f :: * -> *) a. (Foldable f, Eq a) => f a -> [NonEmpty a]
NE.group

{- | /O(n)/

>>> fromDistinctAscListIH [0,1,2]
[0,1,2]
-}
fromDistinctAscListIH :: [Int] -> IntHeap
fromDistinctAscListIH :: [Key] -> IntHeap
fromDistinctAscListIH =
  ([Key] -> IntMap Key) -> [Key] -> IntHeap
forall a b. Coercible a b => a -> b
coerce (forall a. [(Key, a)] -> IntMap a
IM.fromDistinctAscList @Int ([(Key, Key)] -> IntMap Key)
-> ([Key] -> [(Key, Key)]) -> [Key] -> IntMap Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Key -> (Key, Key)) -> [Key] -> [(Key, Key)]
forall a b. (a -> b) -> [a] -> [b]
map ((Key -> Key -> (Key, Key)) -> Key -> Key -> (Key, Key)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) Key
1))

{- | /O(n)/

>>> fromDescListIH [2,1,0,0]
[0,0,1,2]
-}
fromDescListIH :: [Int] -> IntHeap
fromDescListIH :: [Key] -> IntHeap
fromDescListIH =
  IntMap Key -> IntHeap
IntHeap
    (IntMap Key -> IntHeap)
-> ([Key] -> IntMap Key) -> [Key] -> IntHeap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Key, Key)] -> IntMap Key
forall a. [(Key, a)] -> IntMap a
IM.fromDistinctAscList
    ([(Key, Key)] -> IntMap Key)
-> ([Key] -> [(Key, Key)]) -> [Key] -> IntMap Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Key, Key)] -> [(Key, Key)]
forall a. [a] -> [a]
reverse
    ([(Key, Key)] -> [(Key, Key)])
-> ([Key] -> [(Key, Key)]) -> [Key] -> [(Key, Key)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NonEmpty Key -> (Key, Key)) -> [NonEmpty Key] -> [(Key, Key)]
forall a b. (a -> b) -> [a] -> [b]
map (\NonEmpty Key
g -> (NonEmpty Key -> Key
forall a. NonEmpty a -> a
NE.head NonEmpty Key
g, NonEmpty Key -> Key
forall a. NonEmpty a -> Key
NE.length NonEmpty Key
g))
    ([NonEmpty Key] -> [(Key, Key)])
-> ([Key] -> [NonEmpty Key]) -> [Key] -> [(Key, Key)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Key] -> [NonEmpty Key]
forall (f :: * -> *) a. (Foldable f, Eq a) => f a -> [NonEmpty a]
NE.group

{- | /O(n)/

>>> fromDistinctDescListIH [2,1,0]
[0,1,2]
-}
fromDistinctDescListIH :: [Int] -> IntHeap
fromDistinctDescListIH :: [Key] -> IntHeap
fromDistinctDescListIH =
  ([Key] -> IntMap Key) -> [Key] -> IntHeap
forall a b. Coercible a b => a -> b
coerce (forall a. [(Key, a)] -> IntMap a
IM.fromDistinctAscList @Int ([(Key, Key)] -> IntMap Key)
-> ([Key] -> [(Key, Key)]) -> [Key] -> IntMap Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Key -> (Key, Key)) -> [Key] -> [(Key, Key)]
forall a b. (a -> b) -> [a] -> [b]
map ((Key -> Key -> (Key, Key)) -> Key -> Key -> (Key, Key)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) Key
1) ([Key] -> [(Key, Key)])
-> ([Key] -> [Key]) -> [Key] -> [(Key, Key)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Key] -> [Key]
forall a. [a] -> [a]
reverse)

{- | /O(n)/

>>> toListIH (fromListIH [0,1,0,2])
[0,0,1,2]
-}
toListIH :: IntHeap -> [Int]
toListIH :: IntHeap -> [Key]
toListIH = IntHeap -> [Key]
toAscListIH

{- | /O(n)/

>>> toAscListIH (fromListIH [0,1,0,2])
[0,0,1,2]
-}
toAscListIH :: IntHeap -> [Int]
toAscListIH :: IntHeap -> [Key]
toAscListIH =
  forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap @[] (\(Key
k, Key
x) -> Key -> Key -> [Key]
forall a. Key -> a -> [a]
replicate Key
x Key
k)
    ([(Key, Key)] -> [Key])
-> (IntHeap -> [(Key, Key)]) -> IntHeap -> [Key]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IntMap Key -> [(Key, Key)]) -> IntHeap -> [(Key, Key)]
forall a b. Coercible a b => a -> b
coerce (forall a. IntMap a -> [(Key, a)]
IM.toAscList @Int)

{- | /O(n)/

>>> toDescListIH (fromListIH [0,1,0,2])
[2,1,0,0]
-}
toDescListIH :: IntHeap -> [Int]
toDescListIH :: IntHeap -> [Key]
toDescListIH =
  forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap @[] (\(Key
k, Key
x) -> Key -> Key -> [Key]
forall a. Key -> a -> [a]
replicate Key
x Key
k)
    ([(Key, Key)] -> [Key])
-> (IntHeap -> [(Key, Key)]) -> IntHeap -> [Key]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IntMap Key -> [(Key, Key)]) -> IntHeap -> [(Key, Key)]
forall a b. Coercible a b => a -> b
coerce (forall a. IntMap a -> [(Key, a)]
IM.toDescList @Int)

-- | /O(min(n,W))/
insertIH :: Int -> IntHeap -> IntHeap
insertIH :: Key -> IntHeap -> IntHeap
insertIH = (Key -> IntMap Key -> IntMap Key) -> Key -> IntHeap -> IntHeap
forall a b. Coercible a b => a -> b
coerce ((Key -> Key -> IntMap Key -> IntMap Key)
-> Key -> Key -> IntMap Key -> IntMap Key
forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall a. (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
IM.insertWith @Int Key -> Key -> Key
forall a. Num a => a -> a -> a
(+)) Key
1)
{-# INLINE insertIH #-}

{- | /O(min(n,W))/

>>> deleteIH 0 (fromList [0, 0, 1, 2])
[0,1,2]
>>> deleteIH 3 (fromList [0, 0, 1, 2])
[0,0,1,2]
-}
deleteIH :: Int -> IntHeap -> IntHeap
deleteIH :: Key -> IntHeap -> IntHeap
deleteIH = (Key -> IntMap Key -> IntMap Key) -> Key -> IntHeap -> IntHeap
forall a b. Coercible a b => a -> b
coerce (forall a. (a -> Maybe a) -> Key -> IntMap a -> IntMap a
IM.update @Int (\Key
y -> if Key
y Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
1 then Key -> Maybe Key
forall a. a -> Maybe a
Just (Key -> Maybe Key) -> Key -> Maybe Key
forall a b. (a -> b) -> a -> b
$! Key
y Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1 else Maybe Key
forall a. Maybe a
Nothing))

{- | /O(min(n,W))/

>>> deleteAllIH 0 (fromList [0, 0, 1, 2])
[1,2]
>>> deleteAllIH 3 (fromList [0, 0, 1, 2])
[0,0,1,2]
-}
deleteAllIH :: Int -> IntHeap -> IntHeap
deleteAllIH :: Key -> IntHeap -> IntHeap
deleteAllIH = (Key -> IntMap Key -> IntMap Key) -> Key -> IntHeap -> IntHeap
forall a b. Coercible a b => a -> b
coerce (forall a. Key -> IntMap a -> IntMap a
IM.delete @Int)

-- | /O(min(n,W))/
memberIH :: Int -> IntHeap -> Bool
memberIH :: Key -> IntHeap -> Bool
memberIH = (Key -> IntMap Key -> Bool) -> Key -> IntHeap -> Bool
forall a b. Coercible a b => a -> b
coerce (forall a. Key -> IntMap a -> Bool
IM.member @Int)

-- | /O(min(n,W))/
notMemberIH :: Int -> IntHeap -> Bool
notMemberIH :: Key -> IntHeap -> Bool
notMemberIH = (Key -> IntMap Key -> Bool) -> Key -> IntHeap -> Bool
forall a b. Coercible a b => a -> b
coerce (forall a. Key -> IntMap a -> Bool
IM.notMember @Int)

{- | /O(min(n, W))/

>>> countIH 0 (fromList [0, 0, 1, 2])
2
>>> countIH 1 (fromList [0, 0, 1, 2])
1
>>> countIH (-1) (fromList [0, 0, 1, 2])
0
-}
countIH :: Int -> IntHeap -> Int
countIH :: Key -> IntHeap -> Key
countIH = (Key -> IntMap Key -> Key) -> Key -> IntHeap -> Key
forall a b. Coercible a b => a -> b
coerce (forall a. a -> Key -> IntMap a -> a
IM.findWithDefault @Int Key
0)

{- | /O(log n)/

>>> lookupLTIH 1 (fromList [0, 0, 1, 1, 2, 2])
Just 0
>>> lookupLTIH 0 (fromList [0, 0, 1, 1, 2, 2])
Nothing
-}
lookupLTIH :: Int -> IntHeap -> Maybe Int
lookupLTIH :: Key -> IntHeap -> Maybe Key
lookupLTIH Key
x = (IntMap Key -> Maybe Key) -> IntHeap -> Maybe Key
forall a b. Coercible a b => a -> b
coerce (((Key, Key) -> Key) -> Maybe (Key, Key) -> Maybe Key
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Key, Key) -> Key
forall a b. (a, b) -> a
fst (Maybe (Key, Key) -> Maybe Key)
-> (IntMap Key -> Maybe (Key, Key)) -> IntMap Key -> Maybe Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Key -> IntMap a -> Maybe (Key, a)
IM.lookupLT @Int Key
x)

{- | /O(log n)/

>>> lookupGTIH 1 (fromList [0, 0, 1, 1, 2, 2])
Just 2
>>> lookupGTIH 2 (fromList [0, 0, 1, 1, 2, 2])
Nothing
-}
lookupGTIH :: Int -> IntHeap -> Maybe Int
lookupGTIH :: Key -> IntHeap -> Maybe Key
lookupGTIH Key
x = (IntMap Key -> Maybe Key) -> IntHeap -> Maybe Key
forall a b. Coercible a b => a -> b
coerce (((Key, Key) -> Key) -> Maybe (Key, Key) -> Maybe Key
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Key, Key) -> Key
forall a b. (a, b) -> a
fst (Maybe (Key, Key) -> Maybe Key)
-> (IntMap Key -> Maybe (Key, Key)) -> IntMap Key -> Maybe Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Key -> IntMap a -> Maybe (Key, a)
IM.lookupGT @Int Key
x)

{- | /O(log n)/

>>> lookupLEIH 1 (fromList [0, 0, 1, 1, 2, 2])
Just 1
>>> lookupLEIH 0 (fromList [0, 0, 1, 1, 2, 2])
Just 0
>>> lookupLEIH (-1) (fromList [0, 0, 1, 1, 2, 2])
Nothing
-}
lookupLEIH :: Int -> IntHeap -> Maybe Int
lookupLEIH :: Key -> IntHeap -> Maybe Key
lookupLEIH Key
x = (IntMap Key -> Maybe Key) -> IntHeap -> Maybe Key
forall a b. Coercible a b => a -> b
coerce (((Key, Key) -> Key) -> Maybe (Key, Key) -> Maybe Key
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Key, Key) -> Key
forall a b. (a, b) -> a
fst (Maybe (Key, Key) -> Maybe Key)
-> (IntMap Key -> Maybe (Key, Key)) -> IntMap Key -> Maybe Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Key -> IntMap a -> Maybe (Key, a)
IM.lookupLE @Int Key
x)

{- | /O(log n)/

>>> lookupGEIH 1 (fromList [0, 0, 1, 1, 2, 2])
Just 1
>>> lookupGEIH 2 (fromList [0, 0, 1, 1, 2, 2])
Just 2
>>> lookupGEIH 3 (fromList [0, 0, 1, 1, 2, 2])
Nothing
-}
lookupGEIH :: Int -> IntHeap -> Maybe Int
lookupGEIH :: Key -> IntHeap -> Maybe Key
lookupGEIH Key
x = (IntMap Key -> Maybe Key) -> IntHeap -> Maybe Key
forall a b. Coercible a b => a -> b
coerce (((Key, Key) -> Key) -> Maybe (Key, Key) -> Maybe Key
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Key, Key) -> Key
forall a b. (a, b) -> a
fst (Maybe (Key, Key) -> Maybe Key)
-> (IntMap Key -> Maybe (Key, Key)) -> IntMap Key -> Maybe Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Key -> IntMap a -> Maybe (Key, a)
IM.lookupGE @Int Key
x)

-- | /O(1)/
nullIH :: IntHeap -> Bool
nullIH :: IntHeap -> Bool
nullIH = (IntMap Key -> Bool) -> IntHeap -> Bool
forall a b. Coercible a b => a -> b
coerce (forall a. IntMap a -> Bool
IM.null @Int)

{- | /O(n)/

>>> sizeIH (fromList [0, 0, 1, 2])
4
>>> sizeIH emptyIH
0
-}
sizeIH :: IntHeap -> Int
sizeIH :: IntHeap -> Key
sizeIH = (IntMap Key -> Key) -> IntHeap -> Key
forall a b. Coercible a b => a -> b
coerce (forall a b. (a -> b -> a) -> a -> IntMap b -> a
IM.foldl' @Int Key -> Key -> Key
forall a. Num a => a -> a -> a
(+) Key
0)

{- | /O(min(n,W))/

>>> findMinIH (fromList [0, 0, 1, 2])
0
>>> findMinIH emptyIH
*** Exception: findMin: empty map has no minimal element
-}
findMinIH :: IntHeap -> Int
findMinIH :: IntHeap -> Key
findMinIH = (IntMap Key -> Key) -> IntHeap -> Key
forall a b. Coercible a b => a -> b
coerce ((Key, Key) -> Key
forall a b. (a, b) -> a
fst ((Key, Key) -> Key)
-> (IntMap Key -> (Key, Key)) -> IntMap Key -> Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IntMap a -> (Key, a)
IM.findMin @Int)

{- | /O(min(n,W))/

>>> findMaxIH (fromList [0, 0, 1, 2])
2
>>> findMaxIH emptyIH
*** Exception: findMax: empty map has no maximal element
-}
findMaxIH :: IntHeap -> Int
findMaxIH :: IntHeap -> Key
findMaxIH = (IntMap Key -> Key) -> IntHeap -> Key
forall a b. Coercible a b => a -> b
coerce ((Key, Key) -> Key
forall a b. (a, b) -> a
fst ((Key, Key) -> Key)
-> (IntMap Key -> (Key, Key)) -> IntMap Key -> Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IntMap a -> (Key, a)
IM.findMax @Int)

{- | /O(min(n,W))/

>>> deleteMinIH (fromList [0, 0, 1, 2])
[0,1,2]
>>> deleteMinIH emptyIH
*** Exception: updateMinWithKey Nil
-}
deleteMinIH :: IntHeap -> IntHeap
deleteMinIH :: IntHeap -> IntHeap
deleteMinIH =
  (IntMap Key -> IntMap Key) -> IntHeap -> IntHeap
forall a b. Coercible a b => a -> b
coerce (forall a. (a -> Maybe a) -> IntMap a -> IntMap a
IM.updateMin @Int (\Key
x -> if Key
x Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
1 then Key -> Maybe Key
forall a. a -> Maybe a
Just (Key -> Maybe Key) -> Key -> Maybe Key
forall a b. (a -> b) -> a -> b
$! Key
x Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1 else Maybe Key
forall a. Maybe a
Nothing))

{- | /O(min(n,W))/

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

>>> deleteMaxIH emptyIH
*** Exception: updateMaxWithKey Nil
-}
deleteMaxIH :: IntHeap -> IntHeap
deleteMaxIH :: IntHeap -> IntHeap
deleteMaxIH =
  (IntMap Key -> IntMap Key) -> IntHeap -> IntHeap
forall a b. Coercible a b => a -> b
coerce (forall a. (a -> Maybe a) -> IntMap a -> IntMap a
IM.updateMax @Int (\Key
x -> if Key
x Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
1 then Key -> Maybe Key
forall a. a -> Maybe a
Just (Key -> Maybe Key) -> Key -> Maybe Key
forall a b. (a -> b) -> a -> b
$! Key
x Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1 else Maybe Key
forall a. Maybe a
Nothing))

{- | /O(min(n,W))/

>>> deleteFindMinIH (fromList [0, 0, 1, 2])
(0,[0,1,2])
>>> deleteFindMinIH emptyIH
deleteFindMin: empty map has no minimal element
-}
deleteFindMinIH :: IntHeap -> (Int, IntHeap)
deleteFindMinIH :: IntHeap -> (Key, IntHeap)
deleteFindMinIH = (IntMap Key -> (Key, IntMap Key)) -> IntHeap -> (Key, IntHeap)
forall a b. Coercible a b => a -> b
coerce (((Key, Key), IntMap Key) -> (Key, IntMap Key)
found (((Key, Key), IntMap Key) -> (Key, IntMap Key))
-> (IntMap Key -> ((Key, Key), IntMap Key))
-> IntMap Key
-> (Key, IntMap Key)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap Key -> ((Key, Key), IntMap Key)
forall a. IntMap a -> ((Key, a), IntMap a)
IM.deleteFindMin)
  where
    found :: ((Int, Int), IM.IntMap Int) -> (Int, IM.IntMap Int)
    found :: ((Key, Key), IntMap Key) -> (Key, IntMap Key)
found ((Key
k, Key
x), IntMap Key
m)
      | Key
x Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
1 = case Key -> Key -> IntMap Key -> IntMap Key
forall a. Key -> a -> IntMap a -> IntMap a
IM.insert Key
k (Key
x Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1) IntMap Key
m of
          IntMap Key
m' -> (Key
k, IntMap Key
m')
      | Bool
otherwise = (Key
k, IntMap Key
m)

{- | /O(min(n,W))/

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

>>> deleteFindMaxIH emptyIH
deleteFindMax: empty map has no maximal element
-}
deleteFindMaxIH :: IntHeap -> (Int, IntHeap)
deleteFindMaxIH :: IntHeap -> (Key, IntHeap)
deleteFindMaxIH = (IntMap Key -> (Key, IntMap Key)) -> IntHeap -> (Key, IntHeap)
forall a b. Coercible a b => a -> b
coerce (((Key, Key), IntMap Key) -> (Key, IntMap Key)
found (((Key, Key), IntMap Key) -> (Key, IntMap Key))
-> (IntMap Key -> ((Key, Key), IntMap Key))
-> IntMap Key
-> (Key, IntMap Key)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap Key -> ((Key, Key), IntMap Key)
forall a. IntMap a -> ((Key, a), IntMap a)
IM.deleteFindMax)
  where
    found :: ((Int, Int), IM.IntMap Int) -> (Int, IM.IntMap Int)
    found :: ((Key, Key), IntMap Key) -> (Key, IntMap Key)
found ((Key
k, Key
x), IntMap Key
m)
      | Key
x Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
1 = case Key -> Key -> IntMap Key -> IntMap Key
forall a. Key -> a -> IntMap a -> IntMap a
IM.insert Key
k (Key
x Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1) IntMap Key
m of
          IntMap Key
m' -> (Key
k, IntMap Key
m')
      | Bool
otherwise = (Key
k, IntMap Key
m)

{- | /O(min(n,W))/

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

>>> minViewIH emptyIH
Nothing
-}
minViewIH :: IntHeap -> Maybe (Int, IntHeap)
minViewIH :: IntHeap -> Maybe (Key, IntHeap)
minViewIH = (IntMap Key -> Maybe (Key, IntHeap))
-> IntHeap -> Maybe (Key, IntHeap)
forall a b. Coercible a b => a -> b
coerce (Maybe (Key, IntHeap)
-> (((Key, Key), IntMap Key) -> Maybe (Key, IntHeap))
-> Maybe ((Key, Key), IntMap Key)
-> Maybe (Key, IntHeap)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Maybe (Key, IntHeap)
forall a. Maybe a
Nothing ((Key, Key), IntMap Key) -> Maybe (Key, IntHeap)
just (Maybe ((Key, Key), IntMap Key) -> Maybe (Key, IntHeap))
-> (IntMap Key -> Maybe ((Key, Key), IntMap Key))
-> IntMap Key
-> Maybe (Key, IntHeap)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap Key -> Maybe ((Key, Key), IntMap Key)
forall a. IntMap a -> Maybe ((Key, a), IntMap a)
IM.minViewWithKey)
  where
    just :: ((Int, Int), IM.IntMap Int) -> Maybe (Int, IntHeap)
    just :: ((Key, Key), IntMap Key) -> Maybe (Key, IntHeap)
just ((Key
k, Key
x), IntMap Key
m)
      | Key
x Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
1 = case Key -> Key -> IntMap Key -> IntMap Key
forall a. Key -> a -> IntMap a -> IntMap a
IM.insert Key
k (Key
x Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1) IntMap Key
m of
          IntMap Key
m' -> Maybe (Key, IntMap Key) -> Maybe (Key, IntHeap)
forall a b. Coercible a b => a -> b
coerce ((Key, IntMap Key) -> Maybe (Key, IntMap Key)
forall a. a -> Maybe a
Just (Key
k, IntMap Key
m'))
      | Bool
otherwise = Maybe (Key, IntMap Key) -> Maybe (Key, IntHeap)
forall a b. Coercible a b => a -> b
coerce ((Key, IntMap Key) -> Maybe (Key, IntMap Key)
forall a. a -> Maybe a
Just (Key
k, IntMap Key
m))

{- | /O(min(n,W))/

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

>>> maxViewIH emptyIH
Nothing
-}
maxViewIH :: IntHeap -> Maybe (Int, IntHeap)
maxViewIH :: IntHeap -> Maybe (Key, IntHeap)
maxViewIH = (IntMap Key -> Maybe (Key, IntHeap))
-> IntHeap -> Maybe (Key, IntHeap)
forall a b. Coercible a b => a -> b
coerce (Maybe (Key, IntHeap)
-> (((Key, Key), IntMap Key) -> Maybe (Key, IntHeap))
-> Maybe ((Key, Key), IntMap Key)
-> Maybe (Key, IntHeap)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Maybe (Key, IntHeap)
forall a. Maybe a
Nothing ((Key, Key), IntMap Key) -> Maybe (Key, IntHeap)
just (Maybe ((Key, Key), IntMap Key) -> Maybe (Key, IntHeap))
-> (IntMap Key -> Maybe ((Key, Key), IntMap Key))
-> IntMap Key
-> Maybe (Key, IntHeap)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap Key -> Maybe ((Key, Key), IntMap Key)
forall a. IntMap a -> Maybe ((Key, a), IntMap a)
IM.maxViewWithKey)
  where
    just :: ((Int, Int), IM.IntMap Int) -> Maybe (Int, IntHeap)
    just :: ((Key, Key), IntMap Key) -> Maybe (Key, IntHeap)
just ((Key
k, Key
x), IntMap Key
m)
      | Key
x Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
1 = case Key -> Key -> IntMap Key -> IntMap Key
forall a. Key -> a -> IntMap a -> IntMap a
IM.insert Key
k (Key
x Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1) IntMap Key
m of
          IntMap Key
m' -> Maybe (Key, IntMap Key) -> Maybe (Key, IntHeap)
forall a b. Coercible a b => a -> b
coerce ((Key, IntMap Key) -> Maybe (Key, IntMap Key)
forall a. a -> Maybe a
Just (Key
k, IntMap Key
m'))
      | Bool
otherwise = Maybe (Key, IntMap Key) -> Maybe (Key, IntHeap)
forall a b. Coercible a b => a -> b
coerce ((Key, IntMap Key) -> Maybe (Key, IntMap Key)
forall a. a -> Maybe a
Just (Key
k, IntMap Key
m))

{- | /O(min(n,W))/

>>> splitIH 1 (fromList [0, 0, 1, 2])
([0,0],[2])
>>> splitIH 0 (fromList [0, 0, 1, 2])
([],[1,2])
>>> splitIH (-1) (fromList [0, 0, 1, 2])
([],[0,0,1,2])
-}
splitIH :: Int -> IntHeap -> (IntHeap, IntHeap)
splitIH :: Key -> IntHeap -> (IntHeap, IntHeap)
splitIH = (Key -> IntMap Key -> (IntMap Key, IntMap Key))
-> Key -> IntHeap -> (IntHeap, IntHeap)
forall a b. Coercible a b => a -> b
coerce (forall a. Key -> IntMap a -> (IntMap a, IntMap a)
IM.split @Int)

{- | /O(min(n,W))/

>>> splitLookupIH 1 (fromList [0, 0, 1, 2])
([0,0],[1],[2])
>>> splitLookupIH 0 (fromList [0, 0, 1, 2])
([],[0,0],[1,2])
>>> splitLookupIH (-1) (fromList [0, 0, 1, 2])
([],[],[0,0,1,2])
-}
splitLookupIH :: Int -> IntHeap -> (IntHeap, [Int], IntHeap)
splitLookupIH :: Key -> IntHeap -> (IntHeap, [Key], IntHeap)
splitLookupIH Key
k IntHeap
h = case (Key -> IntMap Key -> (IntMap Key, Maybe Key, IntMap Key))
-> Key -> IntHeap -> (IntHeap, Maybe Key, IntHeap)
forall a b. Coercible a b => a -> b
coerce (forall a. Key -> IntMap a -> (IntMap a, Maybe a, IntMap a)
IM.splitLookup @Int) Key
k IntHeap
h of
  (IntHeap
l, Just Key
cnt, IntHeap
r) -> (IntHeap
l, Key -> Key -> [Key]
forall a. Key -> a -> [a]
replicate Key
cnt Key
k, IntHeap
r)
  (IntHeap
l, Maybe Key
Nothing, IntHeap
r) -> (IntHeap
l, [], IntHeap
r)