{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
{- |
= Finite priority heaps
The @'PrioHeap' k a@ type represents a finite heap (or priority queue) from keys/priorities of type @k@ to values of type @a@.
A 'PrioHeap' is strict in its spine. Unlike with maps, duplicate keys/priorities are allowed.
== Performance
The worst case running time complexities are given, with /n/ referring the the number of elements in the heap.
== Warning
The length of a 'PrioHeap' must not exceed @'maxBound' :: 'Int'@.
Violation of this condition is not detected and if the length limit is exceeded, the behaviour of the heap is undefined.
== Implementation
The implementation uses skew binomial heaps, as described by:
* Chris Okasaki, \"Purely Functional Data Structures\", 1998.
-}
module Data.PrioHeap
( PrioHeap
-- * Construction
, empty, singleton
, fromHeap
-- ** From Lists
, fromList
-- * Insertion/Union
, insert
, union, unions
-- * Traversal/Filter
, map, mapWithKey
, traverseWithKey
, filter, filterWithKey
, partition, partitionWithKey
, mapMaybe, mapMaybeWithKey
, mapEither, mapEitherWithKey
-- * Folds
, foldMapWithKey
, foldlWithKey, foldrWithKey
, foldlWithKey', foldrWithKey'
, foldMapOrd
, foldlOrd, foldrOrd
, foldlOrd', foldrOrd'
, foldMapWithKeyOrd
, foldlWithKeyOrd, foldrWithKeyOrd
, foldlWithKeyOrd', foldrWithKeyOrd'
-- * Query
, size
, member, notMember
-- * Min
, adjustMin, adjustMinWithKey
, lookupMin
, findMin
, deleteMin
, deleteFindMin
, updateMin, updateMinWithKey
, minView
-- * Subranges
, take
, drop
, splitAt
, takeWhile, takeWhileWithKey
, dropWhile, dropWhileWithKey
, span, spanWithKey
, break, breakWithKey
, nub
-- * Conversion
, keysHeap
-- ** To Lists
, toList, toAscList, toDescList
) where
import Control.Exception (assert)
import Data.Foldable (foldl', foldr')
import Data.Functor.Classes
import Data.Maybe (fromMaybe)
#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup (Semigroup((<>)))
#endif
import GHC.Exts (IsList)
import qualified GHC.Exts as Exts
import Prelude hiding (break, drop, dropWhile, filter, map, reverse, span, splitAt, take, takeWhile, uncurry)
import Text.Read (Lexeme(Ident), lexP, parens, prec, readPrec)
import Control.DeepSeq (NFData(..))
import qualified Data.Heap.Internal as Heap
import Util.Internal.StrictList
-- | A skew binomial heap with associated priorities.
data PrioHeap k a
= Empty
| Heap
{-# UNPACK #-} !Int -- size
!k -- root key
a -- root value
!(Forest k a) -- forest
type Forest k a = List (Tree k a)
data Pair k a = Pair !k a
data Tree k a = Node
{ _rank :: {-# UNPACK #-} !Int
, _root :: !k
, _value :: a
, _elements :: !(List (Pair k a))
, _children :: !(Forest k a)
}
instance (NFData k, NFData a) => NFData (Pair k a) where
rnf (Pair k x) = rnf k `seq` rnf x
instance (NFData k, NFData a) => NFData (Tree k a) where
rnf (Node _ k x xs c) = rnf k `seq` rnf x `seq` rnf xs `seq` rnf c
errorEmpty :: String -> a
errorEmpty s = error $ "PrioHeap." ++ s ++ ": empty heap"
uncurry :: (a -> b -> c) -> Pair a b -> c
uncurry f (Pair x y) = f x y
{-# INLINE uncurry #-}
link :: Ord k => Tree k a -> Tree k a -> Tree k a
link t1@(Node r1 key1 x1 xs1 c1) t2@(Node r2 key2 x2 xs2 c2) = assert (r1 == r2) $
if key1 <= key2
then Node (r1 + 1) key1 x1 xs1 (t2 `Cons` c1)
else Node (r2 + 1) key2 x2 xs2 (t1 `Cons` c2)
skewLink :: Ord k => k -> a -> Tree k a -> Tree k a -> Tree k a
skewLink kx x t1 t2 = let Node r ky y ys c = link t1 t2
in if kx <= ky
then Node r kx x (Pair ky y `Cons` ys) c
else Node r ky y (Pair kx x `Cons` ys) c
insTree :: Ord k => Tree k a -> Forest k a -> Forest k a
insTree t Nil = t `Cons` Nil
insTree t1 f@(t2 `Cons` ts)
| _rank t1 < _rank t2 = t1 `Cons` f
| otherwise = insTree (link t1 t2) ts
mergeTrees :: Ord k => Forest k a -> Forest k a -> Forest k a
mergeTrees f Nil = f
mergeTrees Nil f = f
mergeTrees f1@(t1 `Cons` ts1) f2@(t2 `Cons` ts2) = case _rank t1 `compare` _rank t2 of
LT -> t1 `Cons` mergeTrees ts1 f2
GT -> t2 `Cons` mergeTrees f1 ts2
EQ -> insTree (link t1 t2) (mergeTrees ts1 ts2)
merge :: Ord k => Forest k a -> Forest k a -> Forest k a
merge f1 f2 = mergeTrees (normalize f1) (normalize f2)
{-# INLINE merge #-}
normalize :: Ord k => Forest k a -> Forest k a
normalize Nil = Nil
normalize (t `Cons` ts) = insTree t ts
{-# INLiNE normalize #-}
ins :: Ord k => k -> a -> Forest k a -> Forest k a
ins key x (t1 `Cons` t2 `Cons` ts)
| _rank t1 == _rank t2 = key `seq` skewLink key x t1 t2 `Cons` ts
ins key x ts = key `seq` Node 0 key x Nil Nil `Cons` ts
fromForest :: Ord k => Int -> Forest k a -> PrioHeap k a
fromForest _ Nil = Empty
fromForest s f@(_ `Cons` _) =
let (Node _ key x xs ts1, ts2) = removeMinTree f
in Heap s key x (foldl' (\acc (Pair key x) -> ins key x acc) (merge (reverse ts1) ts2) xs)
removeMinTree :: Ord k => Forest k a -> (Tree k a, Forest k a)
removeMinTree Nil = error "removeMinTree: empty heap"
removeMinTree (t `Cons` Nil) = (t, Nil)
removeMinTree (t `Cons` ts) =
let (t', ts') = removeMinTree ts
in if _root t <= _root t'
then (t, ts)
else (t', t `Cons` ts')
instance Show2 PrioHeap where
liftShowsPrec2 spk slk spv slv p heap = showsUnaryWith (liftShowsPrec sp sl) "fromList" p (toList heap)
where
sp = liftShowsPrec2 spk slk spv slv
sl = liftShowList2 spk slk spv slv
instance Show k => Show1 (PrioHeap k) where
liftShowsPrec = liftShowsPrec2 showsPrec showList
instance (Show k, Show a) => Show (PrioHeap k a) where
showsPrec = showsPrec2
instance (Ord k, Read k) => Read1 (PrioHeap k) where
liftReadsPrec rp rl = readsData $ readsUnaryWith (liftReadsPrec rp' rl') "fromList" fromList
where
rp' = liftReadsPrec rp rl
rl' = liftReadList rp rl
instance (Ord k, Read k, Read a) => Read (PrioHeap k a) where
readPrec = parens $ prec 10 $ do
Ident "fromList" <- lexP
xs <- readPrec
pure (fromList xs)
instance Ord k => Eq1 (PrioHeap k) where
liftEq f heap1 heap2 = size heap1 == size heap2 && liftEq (liftEq f) (toAscList heap1) (toAscList heap2)
instance (Ord k, Eq a) => Eq (PrioHeap k a) where
(==) = eq1
instance Ord k => Ord1 (PrioHeap k) where
liftCompare f heap1 heap2 = liftCompare (liftCompare f) (toAscList heap1) (toAscList heap2)
instance (Ord k, Ord a) => Ord (PrioHeap k a) where
compare = compare1
instance Ord k => Semigroup (PrioHeap k a) where
(<>) = union
instance Ord k => Monoid (PrioHeap k a) where
mempty = empty
mappend = (<>)
instance Functor (PrioHeap k) where
fmap = map
instance Foldable (PrioHeap k) where
foldMap f = foldMapWithKey (const f)
{-# INLINE foldMap #-}
foldr f = foldrWithKey (const f)
{-# INLINE foldr #-}
foldl f = foldlWithKey (const . f)
{-# INLINE foldl #-}
foldr' f = foldrWithKey' (const f)
{-# INLINE foldr' #-}
foldl' f = foldlWithKey' (const . f)
{-# INLINE foldl' #-}
null Empty = True
null Heap{} = False
length = size
instance Traversable (PrioHeap k) where
traverse f = traverseWithKey (const f)
{-# INLINE traverse #-}
instance Ord k => IsList (PrioHeap k a) where
type Item (PrioHeap k a) = (k, a)
fromList = fromList
toList = toList
instance (NFData k, NFData a) => NFData (PrioHeap k a) where
rnf Empty = ()
rnf (Heap _ k x forest) = rnf k `seq` rnf x `seq` rnf forest
-- | /O(1)/. The empty heap.
--
-- > empty = fromList []
empty :: PrioHeap k a
empty = Empty
-- | /O(1)/. A heap with a single element.
--
-- > singleton x = fromList [x]
singleton :: k -> a -> PrioHeap k a
singleton k x = Heap 1 k x Nil
-- | /O(n * log n)/. Create a heap from a list.
fromList :: Ord k => [(k, a)] -> PrioHeap k a
fromList = foldl' (\acc (key, x) -> insert key x acc) empty
-- | /O(1)/. Insert a new key and value into the heap.
insert :: Ord k => k -> a -> PrioHeap k a -> PrioHeap k a
insert key x Empty = singleton key x
insert kx x (Heap s ky y f)
| kx <= ky = Heap (s + 1) kx x (ins ky y f)
| otherwise = Heap (s + 1) ky y (ins kx x f)
-- | /O(log n)/. The union of two heaps.
union :: Ord k => PrioHeap k a -> PrioHeap k a -> PrioHeap k a
union heap Empty = heap
union Empty heap = heap
union (Heap s1 key1 x1 f1) (Heap s2 key2 x2 f2)
| key1 <= key2 = Heap (s1 + s2) key1 x1 (ins key2 x2 (merge f1 f2))
| otherwise = Heap (s1 + s2) key2 x2 (ins key1 x1 (merge f1 f2))
-- | The union of a foldable of heaps.
--
-- > unions = foldl union empty
unions :: (Foldable f, Ord k) => f (PrioHeap k a) -> PrioHeap k a
unions = foldl' union empty
-- | /O(n)/. Map a function over the heap.
map :: (a -> b) -> PrioHeap k a -> PrioHeap k b
map f = mapWithKey (const f)
{-# INLINE map #-}
-- | /O(n)/. Map a function that has access to the key associated with a value over the heap.
mapWithKey :: (k -> a -> b) -> PrioHeap k a -> PrioHeap k b
mapWithKey _ Empty = Empty
mapWithKey f (Heap s key x forest) = Heap s key (f key x) (fmap mapTree forest)
where
mapTree (Node r key x xs c) = Node r key (f key x) (fmap mapPair xs) (fmap mapTree c)
mapPair (Pair key x) = Pair key (f key x)
-- | /O(n)/. Traverse the heap with a function that has access to the key associated with a value.
traverseWithKey :: Applicative f => (k -> a -> f b) -> PrioHeap k a -> f (PrioHeap k b)
traverseWithKey f = go
where
go Empty = pure Empty
go (Heap s key x forest) = Heap s key <$> f key x <*> traverse traverseTree forest
traverseTree (Node r key x xs c) = Node r key <$> f key x <*> traverse traversePair xs <*> traverse traverseTree c
traversePair (Pair key x) = Pair key <$> f key x
{-# INLINE traverseWithKey #-}
-- | /O(n)/. Filter all elements that satisfy the predicate.
filter :: Ord k => (a -> Bool) -> PrioHeap k a -> PrioHeap k a
filter f = filterWithKey (const f)
{-# INLINE filter #-}
-- | /O(n)/. Filter all elements that satisfy the predicate.
filterWithKey :: Ord k => (k -> a -> Bool) -> PrioHeap k a -> PrioHeap k a
filterWithKey f = foldrWithKey f' empty
where
f' key x heap
| f key x = insert key x heap
| otherwise = heap
{-# INLINE filterWithKey #-}
-- | /O(n)/. Partition the heap into two heaps, one with all elements that satisfy the predicate
-- and one with all elements that don't satisfy the predicate.
partition :: Ord k => (a -> Bool) -> PrioHeap k a -> (PrioHeap k a, PrioHeap k a)
partition f = partitionWithKey (const f)
{-# INLINE partition #-}
-- | /O(n)/. Partition the heap into two heaps, one with all elements that satisfy the predicate
-- and one with all elements that don't satisfy the predicate.
partitionWithKey :: Ord k => (k -> a -> Bool) -> PrioHeap k a -> (PrioHeap k a, PrioHeap k a)
partitionWithKey f = foldrWithKey f' (empty, empty)
where
f' key x (heap1, heap2)
| f key x = (insert key x heap1, heap2)
| otherwise = (heap1, insert key x heap2)
{-# INLINE partitionWithKey #-}
-- | /O(n)/. Map and collect the 'Just' results.
mapMaybe :: Ord k => (a -> Maybe b) -> PrioHeap k a -> PrioHeap k b
mapMaybe f = mapMaybeWithKey (const f)
{-# INLINE mapMaybe #-}
-- | /O(n)/. Map and collect the 'Just' results.
mapMaybeWithKey :: Ord k => (k -> a -> Maybe b) -> PrioHeap k a -> PrioHeap k b
mapMaybeWithKey f = foldrWithKey f' empty
where
f' key x heap = case f key x of
Just y -> insert key y heap
Nothing -> heap
{-# INLINE mapMaybeWithKey #-}
-- | /O(n)/. Map and separate the 'Left' and 'Right' results.
mapEither :: Ord k => (a -> Either b c) -> PrioHeap k a -> (PrioHeap k b, PrioHeap k c)
mapEither f = mapEitherWithKey (const f)
{-# INLINE mapEither #-}
-- | /O(n)/. Map and separate the 'Left' and 'Right' results.
mapEitherWithKey :: Ord k => (k -> a -> Either b c) -> PrioHeap k a -> (PrioHeap k b, PrioHeap k c)
mapEitherWithKey f = foldrWithKey f' (empty, empty)
where
f' key x (heap1, heap2) = case f key x of
Left y -> (insert key y heap1, heap2)
Right y -> (heap1, insert key y heap2)
{-# INLINE mapEitherWithKey #-}
-- | /O(n)/. Fold the keys and values in the heap, using the given monoid.
foldMapWithKey :: Monoid m => (k -> a -> m) -> PrioHeap k a -> m
foldMapWithKey f = foldrWithKey (\key x acc -> f key x `mappend` acc) mempty
{-# INLINE foldMapWithKey #-}
-- | /O(n)/. Fold the keys and values in the heap, using the given right-associative function.
foldrWithKey :: (k -> a -> b -> b) -> b -> PrioHeap k a -> b
foldrWithKey f acc = go
where
go Empty = acc
go (Heap _ key x forest) = f key x (foldr foldTree acc forest)
foldTree (Node _ key x xs c) acc = f key x (foldr (uncurry f) (foldr foldTree acc c) xs)
{-# INLINE foldrWithKey #-}
-- | /O(n)/. Fold the keys and values in the heap, using the given left-associative function.
foldlWithKey :: (b -> k -> a -> b) -> b -> PrioHeap k a -> b
foldlWithKey f acc = go
where
go Empty = acc
go (Heap _ key x forest) = foldl foldTree (f acc key x) forest
foldTree acc (Node _ key x xs c) = foldl foldTree (foldl (uncurry . f) (f acc key x) xs) c
{-# INLINE foldlWithKey #-}
-- | /O(n)/. A strict version of 'foldrWithKey'.
-- Each application of the function is evaluated before using the result in the next application.
foldrWithKey' :: (k -> a -> b -> b) -> b -> PrioHeap k a -> b
foldrWithKey' f acc h = foldlWithKey f' id h acc
where
f' k key x z = k $! f key x z
{-# INLINE foldrWithKey' #-}
-- | /O(n)/. A strict version of 'foldlWithKey'.
-- Each application of the function is evaluated before using the result in the next application.
foldlWithKey' :: (b -> k -> a -> b) -> b -> PrioHeap k a -> b
foldlWithKey' f acc h = foldrWithKey f' id h acc
where
f' key x k z = k $! f z key x
{-# INLINE foldlWithKey' #-}
-- | /O(n * log n)/. Fold the values in the heap in order, using the given monoid.
foldMapOrd :: (Ord k, Monoid m) => (a -> m) -> PrioHeap k a -> m
foldMapOrd f = foldMapWithKeyOrd (const f)
{-# INLINE foldMapOrd #-}
-- | /O(n * log n)/. Fold the values in the heap in order, using the given right-associative function.
foldrOrd :: Ord k => (a -> b -> b) -> b -> PrioHeap k a -> b
foldrOrd f = foldrWithKeyOrd (const f)
{-# INLINE foldrOrd #-}
-- | /O(n * log n)/. Fold the values in the heap in order, using the given left-associative function.
foldlOrd :: Ord k => (b -> a -> b) -> b -> PrioHeap k a -> b
foldlOrd f = foldlWithKeyOrd (const . f)
{-# INLINE foldlOrd #-}
-- | /O(n * log n)/. A strict version of 'foldrOrd'.
-- Each application of the function is evaluated before using the result in the next application.
foldrOrd' :: Ord k => (a -> b -> b) -> b -> PrioHeap k a -> b
foldrOrd' f = foldrWithKeyOrd' (const f)
{-# INLINE foldrOrd' #-}
-- | /O(n)/. A strict version of 'foldlOrd'.
-- Each application of the function is evaluated before using the result in the next application.
foldlOrd' :: Ord k => (b -> a -> b) -> b -> PrioHeap k a -> b
foldlOrd' f = foldlWithKeyOrd' (const . f)
{-# INLINE foldlOrd' #-}
-- | /O(n * log n)/. Fold the keys and values in the heap in order, using the given monoid.
foldMapWithKeyOrd :: (Ord k, Monoid m) => (k -> a -> m) -> PrioHeap k a -> m
foldMapWithKeyOrd f = foldrWithKeyOrd (\key x acc -> f key x `mappend` acc) mempty
{-# INLINE foldMapWithKeyOrd #-}
-- | /O(n * log n)/. Fold the keys and values in the heap in order, using the given right-associative function.
foldrWithKeyOrd :: Ord k => (k -> a -> b -> b) -> b -> PrioHeap k a -> b
foldrWithKeyOrd f acc = go
where
go h = case minView h of
Nothing -> acc
Just ((key, x), h') -> f key x (go h')
{-# INLINE foldrWithKeyOrd #-}
-- | /O(n * log n)/. Fold the keys and values in the heap in order, using the given left-associative function.
foldlWithKeyOrd :: Ord k => (b -> k -> a -> b) -> b -> PrioHeap k a -> b
foldlWithKeyOrd f = go
where
go acc h = case minView h of
Nothing -> acc
Just ((key, x), h') -> go (f acc key x) h'
{-# INLINE foldlWithKeyOrd #-}
-- | /O(n * log n)/. A strict version of 'foldrWithKeyOrd'.
-- Each application of the function is evaluated before using the result in the next application.
foldrWithKeyOrd' :: Ord k => (k -> a -> b -> b) -> b -> PrioHeap k a -> b
foldrWithKeyOrd' f acc h = foldlWithKeyOrd f' id h acc
where
f' k key x z = k $! f key x z
{-# INLINE foldrWithKeyOrd' #-}
-- | /O(n)/. A strict version of 'foldlWithKeyOrd'.
-- Each application of the function is evaluated before using the result in the next application.
foldlWithKeyOrd' :: Ord k => (b -> k -> a -> b) -> b -> PrioHeap k a -> b
foldlWithKeyOrd' f acc h = foldrWithKeyOrd f' id h acc
where
f' key x k z = k $! f z key x
{-# INLINE foldlWithKeyOrd' #-}
-- | /O(1)/. The number of elements in the heap.
size :: PrioHeap k a -> Int
size Empty = 0
size (Heap s _ _ _) = s
-- | /O(n)/. Is the key a member of the heap?
member :: Ord k => k -> PrioHeap k a -> Bool
member _ Empty = False
member kx (Heap _ ky _ forest) = kx <= ky && any (kx `elemTree`) forest
where
kx `elemTree` (Node _ ky _ ys c) = kx <= ky && (any (\(Pair a _) -> kx == a) ys || any (kx `elemTree`) c)
-- | /O(n)/. Is the key not a member of the heap?
notMember :: Ord k => k -> PrioHeap k a -> Bool
notMember key = not . member key
-- | /O(1)/. Adjust the value at the minimal key.
adjustMin :: (a -> a) -> PrioHeap k a -> PrioHeap k a
adjustMin f = adjustMinWithKey (const f)
{-# INLINE adjustMin #-}
-- | /O(1)/. Adjust the value at the minimal key.
adjustMinWithKey :: (k -> a -> a) -> PrioHeap k a -> PrioHeap k a
adjustMinWithKey _ Empty = Empty
adjustMinWithKey f (Heap s key x forest) = Heap s key (f key x) forest
-- | /O(1)/. The minimal element in the heap or 'Nothing' if the heap is empty.
lookupMin :: PrioHeap k a -> Maybe (k, a)
lookupMin Empty = Nothing
lookupMin (Heap _ key x _) = Just (key, x)
-- | /O(1)/. The minimal element in the heap. Calls 'error' if the heap is empty.
findMin :: PrioHeap k a -> (k, a)
findMin heap = fromMaybe (errorEmpty "findMin") (lookupMin heap)
-- | /O(log n)/. Delete the minimal element. Returns the empty heap if the heap is empty.
deleteMin :: Ord k => PrioHeap k a -> PrioHeap k a
deleteMin Empty = Empty
deleteMin (Heap s _ _ f) = fromForest (s - 1) f
-- | /O(log n)/. Delete and find the minimal element. Calls 'error' if the heap is empty.
--
-- > deleteFindMin heap = (findMin heap, deleteMin heap)
deleteFindMin :: Ord k => PrioHeap k a -> ((k, a), PrioHeap k a)
deleteFindMin heap = fromMaybe (errorEmpty "deleteFindMin") (minView heap)
-- | /O(log n)/. Update the value at the minimal key.
updateMin :: Ord k => (a -> Maybe a) -> PrioHeap k a -> PrioHeap k a
updateMin f = updateMinWithKey (const f)
{-# INLINE updateMin #-}
-- | /O(log n)/. Update the value at the minimal key.
updateMinWithKey :: Ord k => (k -> a -> Maybe a) -> PrioHeap k a -> PrioHeap k a
updateMinWithKey _ Empty = Empty
updateMinWithKey f (Heap s key x forest) = case f key x of
Nothing -> fromForest (s - 1) forest
Just x' -> Heap s key x' forest
-- | /O(log n)/. Retrieves the minimal key/value pair of the heap and the heap stripped of that element or 'Nothing' if the heap is empty.
minView :: Ord k => PrioHeap k a -> Maybe ((k, a), PrioHeap k a)
minView Empty = Nothing
minView (Heap s key x f) = Just ((key, x), fromForest (s - 1) f)
-- | /O(n * log n)/. @take n heap@ takes the @n@ smallest elements of @heap@, in ascending order.
--
-- > take n heap = take n (toAscList heap)
take :: Ord k => Int -> PrioHeap k a -> [(k, a)]
take n h
| n <= 0 = []
| otherwise = case minView h of
Nothing -> []
Just (x, h') -> x : take (n - 1) h'
-- | /O(n * log n)/. @drop n heap@ drops the @n@ smallest elements from @heap@.
drop :: Ord k => Int -> PrioHeap k a -> PrioHeap k a
drop n h
| n <= 0 = h
| otherwise = drop (n - 1) (deleteMin h)
-- | /O(n * log n)/. @splitAt n heap@ takes and drops the @n@ smallest elements from @heap@.
splitAt :: Ord k => Int -> PrioHeap k a -> ([(k, a)], PrioHeap k a)
splitAt n h
| n <= 0 = ([], h)
| otherwise = case minView h of
Nothing -> ([], h)
Just (x, h') -> let (xs, h'') = splitAt (n - 1) h' in (x : xs, h'')
-- | /O(n * log n)/. @takeWhile p heap@ takes the elements from @heap@ in ascending order, while @p@ holds.
takeWhile :: Ord k => (a -> Bool) -> PrioHeap k a -> [(k, a)]
takeWhile p = takeWhileWithKey (const p)
{-# INLINE takeWhile #-}
-- | /O(n * log n)/. @takeWhileWithKey p heap@ takes the elements from @heap@ in ascending order, while @p@ holds.
takeWhileWithKey :: Ord k => (k -> a -> Bool) -> PrioHeap k a -> [(k, a)]
takeWhileWithKey p = go
where
go h = case minView h of
Nothing -> []
Just ((key, x), h') -> if p key x then (key, x) : go h' else []
{-# INLINE takeWhileWithKey #-}
-- | /O(n * log n)/. @dropWhile p heap@ drops the elements from @heap@ in ascending order, while @p@ holds.
dropWhile :: Ord k => (a -> Bool) -> PrioHeap k a -> PrioHeap k a
dropWhile p = dropWhileWithKey (const p)
{-# INLINE dropWhile #-}
-- | /O(n * log n)/. @dropWhileWithKey p heap@ drops the elements from @heap@ in ascending order, while @p@ holds.
dropWhileWithKey :: Ord k => (k -> a -> Bool) -> PrioHeap k a -> PrioHeap k a
dropWhileWithKey p = go
where
go h = case minView h of
Nothing -> h
Just ((key, x), h') -> if p key x then go h' else h
{-# INLINE dropWhileWithKey #-}
-- | /O(n * log n)/. @span p heap@ takes and drops the elements from @heap@, while @p@ holds
span :: Ord k => (a -> Bool) -> PrioHeap k a -> ([(k, a)], PrioHeap k a)
span p = spanWithKey (const p)
{-# INLINE span #-}
-- | /O(n * log n)/. @spanWithKey p heap@ takes and drops the elements from @heap@, while @p@ holds
spanWithKey :: Ord k => (k -> a -> Bool) -> PrioHeap k a -> ([(k, a)], PrioHeap k a)
spanWithKey p = go
where
go h = case minView h of
Nothing -> ([], h)
Just ((key, x), h') -> if p key x
then let (xs, h'') = go h' in ((key, x) : xs, h'')
else ([], h)
{-# INLINE spanWithKey #-}
-- | /O(n * log n)/. @span@, but with inverted predicate.
break :: Ord k => (a -> Bool) -> PrioHeap k a -> ([(k, a)], PrioHeap k a)
break p = span (not . p)
{-# INLINE break #-}
-- | /O(n * log n)/. @spanWithKey@, but with inverted predicate.
breakWithKey :: Ord k => (k -> a -> Bool) -> PrioHeap k a -> ([(k, a)], PrioHeap k a)
breakWithKey p = spanWithKey (\key x -> not (p key x))
{-# INLINE breakWithKey #-}
-- | /O(n * log n)/. Remove duplicate elements from the heap.
nub :: Ord k => PrioHeap k a -> PrioHeap k a
nub h = case minView h of
Nothing -> Empty
Just ((key, x), h') -> insert key x (nub (dropWhileWithKey (const . (== key)) h'))
-- | /O(n)/. Create a list of key/value pairs from the heap.
toList :: PrioHeap k a -> [(k, a)]
toList = foldrWithKey (\key x acc -> (key, x) : acc) []
-- | /O(n * log n)/. Create an ascending list of key/value pairs from the heap.
toAscList :: Ord k => PrioHeap k a -> [(k, a)]
toAscList = foldrWithKeyOrd (\key x acc -> (key, x) : acc) []
-- | /O(n * log n)/. Create a descending list of key/value pairs from the heap.
toDescList :: Ord k => PrioHeap k a -> [(k, a)]
toDescList = foldlWithKeyOrd (\acc key x -> (key, x) : acc) []
-- | /O(n)/. Create a heap from a 'Data.Heap.Heap' of keys and a function which computes the value for each key.
fromHeap :: (k -> a) -> Heap.Heap k -> PrioHeap k a
fromHeap _ Heap.Empty = Empty
fromHeap f (Heap.Heap s key forest) = Heap s key (f key) (fmap fromTree forest)
where
fromTree (Heap.Node r key xs c) = Node r key (f key) (fmap (\key -> Pair key (f key)) xs) (fmap fromTree c)
-- | Create a 'Data.Heap.Heap' of all keys of the heap
keysHeap :: PrioHeap k a -> Heap.Heap k
keysHeap Empty = Heap.Empty
keysHeap (Heap s key _ forest) = Heap.Heap s key (fmap fromTree forest)
where
fromTree (Node r key _ xs c) = Heap.Node r key (fmap (\(Pair key _) -> key) xs) (fmap fromTree c)