{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Heap.Internal
( Heap(..)
, Tree(..)
-- * Construction
, empty, singleton
-- ** From Lists
, fromList
-- * Insertion/Union
, insert
, union, unions
-- * Traversal/Filter
, map, mapMonotonic
, filter
, partition
-- * Ordered Folds
, foldMapOrd
, foldlOrd, foldrOrd
, foldlOrd', foldrOrd'
-- * Query
, size
, member, notMember
-- * Min
, lookupMin
, findMin
, deleteMin
, deleteFindMin
, minView
-- * Subranges
, take
, drop
, splitAt
, takeWhile
, dropWhile
, span
, break
, nub
-- * Conversion
-- ** To Lists
, toAscList, toDescList
-- * Heapsort
, heapsort
) where
import Control.Exception (assert)
import Data.Foldable (foldl', toList)
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)
import Text.Read (Lexeme(Ident), lexP, parens, prec, readPrec)
import Control.DeepSeq (NFData(..))
import Util.Internal.StrictList
-- | A skew binomial heap.
data Heap a
= Empty
| Heap
{-# UNPACK #-} !Int -- size
!a -- root
!(Forest a) -- forest
type Forest a = List (Tree a)
data Tree a = Node
{ _rank :: {-# UNPACK #-} !Int
, _root :: !a
, _elements :: !(List a)
, _children :: !(Forest a)
}
instance NFData a => NFData (Tree a) where
rnf (Node _ x xs c) = rnf x `seq` rnf xs `seq` rnf c
errorEmpty :: String -> a
errorEmpty s = error $ "Heap." ++ s ++ ": empty heap"
link :: Ord a => Tree a -> Tree a -> Tree a
link t1@(Node r1 x1 xs1 c1) t2@(Node r2 x2 xs2 c2) = assert (r1 == r2) $
if x1 <= x2
then Node (r1 + 1) x1 xs1 (t2 `Cons` c1)
else Node (r2 + 1) x2 xs2 (t1 `Cons` c2)
skewLink :: Ord a => a -> Tree a -> Tree a -> Tree a
skewLink x t1 t2 = let Node r y ys c = link t1 t2
in if x <= y
then Node r x (y `Cons` ys) c
else Node r y (x `Cons` ys) c
insTree :: Ord a => Tree a -> Forest a -> Forest 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 a => Forest a -> Forest a -> Forest 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 a => Forest a -> Forest a -> Forest a
merge f1 f2 = mergeTrees (normalize f1) (normalize f2)
{-# INLINE merge #-}
normalize :: Ord a => Forest a -> Forest a
normalize Nil = Nil
normalize (t `Cons` ts) = insTree t ts
{-# INLiNE normalize #-}
ins :: Ord a => a -> Forest a -> Forest a
ins x (t1 `Cons` t2 `Cons` ts)
| _rank t1 == _rank t2 = x `seq` skewLink x t1 t2 `Cons` ts
ins x ts = x `seq` Node 0 x Nil Nil `Cons` ts
fromForest :: Ord a => Int -> Forest a -> Heap a
fromForest _ Nil = Empty
fromForest s f@(_ `Cons` _) =
let (Node _ x xs ts1, ts2) = removeMinTree f
in Heap s x (foldl' (flip ins) (merge (reverse ts1) ts2) xs)
removeMinTree :: Ord a => Forest a -> (Tree a, Forest 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 Show1 Heap where
liftShowsPrec sp sl p heap = showsUnaryWith (liftShowsPrec sp sl) "fromList" p (toList heap)
instance Show a => Show (Heap a) where
showsPrec = showsPrec1
instance (Ord a, Read a) => Read (Heap a) where
readPrec = parens $ prec 10 $ do
Ident "fromList" <- lexP
xs <- readPrec
pure (fromList xs)
instance Ord a => Eq (Heap a) where
heap1 == heap2 = size heap1 == size heap2 && toAscList heap1 == toAscList heap2
instance Ord a => Ord (Heap a) where
compare heap1 heap2 = compare (toAscList heap1) (toAscList heap2)
instance Ord a => Semigroup (Heap a) where
(<>) = union
instance Ord a => Monoid (Heap a) where
mempty = empty
mappend = (<>)
instance Foldable Heap where
foldr f acc = go
where
go Empty = acc
go (Heap _ x forest) = f x (foldr foldTree acc forest)
foldTree (Node _ x xs c) acc = f x (foldr f (foldr foldTree acc c) xs)
{-# INLINE foldr #-}
foldl f acc = go
where
go Empty = acc
go (Heap _ x forest) = foldl foldTree (f acc x) forest
foldTree acc (Node _ x xs c) = foldl foldTree (foldl f (f acc x) xs) c
{-# INLINE foldl #-}
null Empty = True
null Heap{} = False
length = size
minimum = findMin
instance Ord a => IsList (Heap a) where
type Item (Heap a) = a
fromList = fromList
toList = toList
instance NFData a => NFData (Heap a) where
rnf Empty = ()
rnf (Heap _ x forest) = rnf x `seq` rnf forest
-- | /O(1)/. The empty heap.
--
-- > empty = fromList []
empty :: Heap a
empty = Empty
-- | /O(1)/. A heap with a single element.
--
-- > singleton x = fromList [x]
singleton :: a -> Heap a
singleton x = Heap 1 x Nil
-- | /O(n)/. Create a heap from a list.
fromList :: Ord a => [a] -> Heap a
fromList = foldl' (flip insert) empty
-- | /O(1)/. Insert a new value into the heap.
insert :: Ord a => a -> Heap a -> Heap a
insert x Empty = singleton x
insert x (Heap s y f)
| x <= y = Heap (s + 1) x (ins y f)
| otherwise = Heap (s + 1) y (ins x f)
-- | /O(log n)/. The union of two heaps.
union :: Ord a => Heap a -> Heap a -> Heap a
union heap Empty = heap
union Empty heap = heap
union (Heap s1 x1 f1) (Heap s2 x2 f2)
| x1 <= x2 = Heap (s1 + s2) x1 (ins x2 (merge f1 f2))
| otherwise = Heap (s1 + s2) x2 (ins x1 (merge f1 f2))
-- | The union of a foldable of heaps.
--
-- > unions = foldl union empty
unions :: (Foldable f, Ord a) => f (Heap a) -> Heap a
unions = foldl' union empty
-- | /O(n)/. Map a function over the heap.
map :: Ord b => (a -> b) -> Heap a -> Heap b
map f = fromList . fmap f . toList
-- | /O(n)/, Map an increasing function over the heap. The precondition is not checked.
mapMonotonic :: (a -> b) -> Heap a -> Heap b
mapMonotonic _ Empty = Empty
mapMonotonic f (Heap s x forest) = Heap s (f x) (fmap mapTree forest)
where
mapTree (Node r x xs c) = Node r (f x) (fmap f xs) (fmap mapTree c)
-- | /O(n)/. Filter all elements that satisfy the predicate.
filter :: Ord a => (a -> Bool) -> Heap a -> Heap a
filter f = foldl' (\acc x -> if f x then insert x acc else acc) empty
-- | /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 a => (a -> Bool) -> Heap a -> (Heap a, Heap a)
partition f = foldl' (\(h1, h2) x -> if f x then (insert x h1, h2) else (h1, insert x h2)) (empty, empty)
-- | /O(n * log n)/. Fold the values in the heap in order, using the given monoid.
foldMapOrd :: (Ord a, Monoid m) => (a -> m) -> Heap a -> m
foldMapOrd f = foldrOrd (mappend . f) mempty
-- | /O(n * log n)/. Fold the values in the heap in order, using the given right-associative function.
foldrOrd :: Ord a => (a -> b -> b) -> b -> Heap a -> b
foldrOrd f acc = go
where
go h = case minView h of
Nothing -> acc
Just (x, h') -> f x (go h')
{-# INLINE foldrOrd #-}
-- | /O(n * log n)/. Fold the values in the heap in order, using the given left-associative function.
foldlOrd :: Ord a => (b -> a -> b) -> b -> Heap a -> b
foldlOrd f = go
where
go acc h = case minView h of
Nothing -> acc
Just (x, h') -> go (f acc x) h'
{-# 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 a => (a -> b -> b) -> b -> Heap a -> b
foldrOrd' f acc h = foldlOrd f' id h acc
where
f' k x z = k $! f x z
{-# 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 a => (b -> a -> b) -> b -> Heap a -> b
foldlOrd' f acc h = foldrOrd f' id h acc
where
f' x k z = k $! f z x
{-# INLINE foldlOrd' #-}
-- | /O(1)/. The number of elements in the heap.
size :: Heap a -> Int
size Empty = 0
size (Heap s _ _) = s
-- | /O(n)/. Is the value a member of the heap?
member :: Ord a => a -> Heap a -> Bool
member _ Empty = False
member x (Heap _ y forest) = x <= y && any (x `elemTree`) forest
where
x `elemTree` (Node _ y ys c) = x <= y && (x `elem` ys || any (x `elemTree`) c)
-- | /O(n)/. Is the value not a member of the heap?
notMember :: Ord a => a -> Heap a -> Bool
notMember x = not . member x
-- | /O(log n)/. The minimal element in the heap. Calls 'error' if the heap is empty.
findMin :: Heap a -> a
findMin heap = fromMaybe (errorEmpty "findMin") (lookupMin heap)
-- | /O(log n)/. The minimal element in the heap or 'Nothing' if the heap is empty.
lookupMin :: Heap a -> Maybe a
lookupMin Empty = Nothing
lookupMin (Heap _ x _) = Just $! x
-- | /O(log n)/. Delete the minimal element. Returns the empty heap if the heap is empty.
deleteMin :: Ord a => Heap a -> Heap 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 a => Heap a -> (a, Heap a)
deleteFindMin heap = fromMaybe (errorEmpty "deleteFindMin") (minView heap)
-- | /O(log n)/. Retrieves the minimal element of the heap and the heap stripped of that element or 'Nothing' if the heap is empty.
minView :: Ord a => Heap a -> Maybe (a, Heap a)
minView Empty = Nothing
minView (Heap s x f) = Just (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 a => Int -> Heap a -> [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 a => Int -> Heap a -> Heap 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 n heap = (take n heap, drop n heap)
splitAt :: Ord a => Int -> Heap a -> ([a], Heap 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 a => (a -> Bool) -> Heap a -> [a]
takeWhile p = go
where
go h = case minView h of
Nothing -> []
Just (x, h') -> if p x then x : go h' else []
{-# INLINE takeWhile #-}
-- | /O(n * log n)/. @dropWhile p heap@ drops the elements from @heap@ in ascending order, while @p@ holds.
dropWhile :: Ord a => (a -> Bool) -> Heap a -> Heap a
dropWhile p = go
where
go h = case minView h of
Nothing -> h
Just (x, h') -> if p x then go h' else h
{-# INLINE dropWhile #-}
-- | /O(n * log n)/. @span p heap@ takes and drops the elements from @heap@, while @p@ holds
--
-- > span p heap = (takeWhile p heap, dropWhile p heap)
span :: Ord a => (a -> Bool) -> Heap a -> ([a], Heap a)
span p = go
where
go h = case minView h of
Nothing -> ([], h)
Just (x, h') -> if p x
then let (xs, h'') = go h' in (x : xs, h'')
else ([], h)
{-# INLINE span #-}
-- | /O(n * log n)/. @span@, but with inverted predicate.
--
-- > break p = span (not . p)
break :: Ord a => (a -> Bool) -> Heap a -> ([a], Heap a)
break p = span (not . p)
{-# INLINE break #-}
-- | /O(n * log n)/. Remove duplicate elements from the heap.
nub :: Ord a => Heap a -> Heap a
nub h = case minView h of
Nothing -> Empty
Just (x, h') -> insert x (nub (dropWhile (== x) h'))
-- | /O(n * log n)/. Create a descending list from the heap.
toAscList :: Ord a => Heap a -> [a]
toAscList = foldrOrd (:) []
-- | /O(n * log n)/. Create a descending list from the heap.
toDescList :: Ord a => Heap a -> [a]
toDescList = foldlOrd (flip (:)) []
-- | /O(n * log n)/. Sort a list using a heap. The sort is unstable.
heapsort :: Ord a => [a] -> [a]
heapsort = toAscList . fromList