{-# LANGUAGE CPP #-}
#ifdef __GLASGOW_HASKELL__
{-# LANGUAGE TypeFamilies #-}
#endif
{- |
= 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 in
* 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
#ifdef __GLASGOW_HASKELL__
import GHC.Exts (IsList)
import qualified GHC.Exts as Exts
#endif
import Prelude hiding (break, drop, dropWhile, filter, map, reverse, span, splitAt, take, takeWhile, uncurry)
import Text.Read (Lexeme(Ident), lexP, parens, prec, readPrec)
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)
}
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
{-# INLINE liftShowsPrec #-}
instance (Show k, Show a) => Show (PrioHeap k a) where
showsPrec = showsPrec2
{-# INLINE showsPrec #-}
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
#ifdef __GLASGOW_HASKELL__
readPrec = parens $ prec 10 $ do
Ident "fromList" <- lexP
xs <- readPrec
pure (fromList xs)
#else
readsPrec = readsPrec1
{-# INLINE readPrec #-}
#endif
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
{-# INLINE (==) #-}
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
{-# INLINE compare #-}
instance Ord k => Semigroup (PrioHeap k a) where
(<>) = union
{-# INLINE (<>) #-}
instance Ord k => Monoid (PrioHeap k a) where
mempty = empty
{-# INLINE mempty #-}
mappend = (<>)
{-# INLINE mappend #-}
instance Functor (PrioHeap k) where
fmap = map
{-# INLINE fmap #-}
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
{-# INLINE null #-}
length = size
{-# INLINE length #-}
instance Traversable (PrioHeap k) where
traverse f = traverseWithKey (const f)
{-# INLINE traverse #-}
#ifdef __GLASGOW_HASKELL__
instance Ord k => IsList (PrioHeap k a) where
type Item (PrioHeap k a) = (k, a)
fromList = fromList
{-# INLINE fromList #-}
toList = toList
{-# INLINE toList #-}
#endif
-- | /O(1)/. The empty heap.
--
-- > empty = fromList []
empty :: PrioHeap k a
empty = Empty
{-# INLINE 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
{-# INLINE singleton #-}
-- | /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
{-# INLINE fromList #-}
-- | /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
{-# INLINE unions #-}
-- | /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)
{-# INLINE mapWithKey #-}
-- | /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 _ Empty = pure Empty
traverseWithKey f (Heap s key x forest) = Heap s key <$> f key x <*> traverse traverseTree forest
where
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 _ acc Empty = acc
foldrWithKey f acc (Heap _ key x forest) = f key x (foldr foldTree acc forest)
where
foldTree (Node _ key x xs c) acc = f key x (foldr (uncurry f) (foldr foldTree acc c) xs)
-- | /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 _ acc Empty = acc
foldlWithKey f acc (Heap _ key x forest) = foldl foldTree (f acc key x) forest
where
foldTree acc (Node _ key x xs c) = foldl foldTree (foldl (uncurry . f) (f acc key x) xs) c
-- | /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
{-# INLINE size #-}
-- | /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 value 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)
{-# INLINE lookupMin #-}
-- | /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)
{-# INLINE findMin #-}
-- | /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)
{-# INLINE deleteFindMin #-}
-- | /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)
{-# INLINE minView #-}
-- | /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)