module Data.Heap(
Heap, empty,
fromList, toList, elems,
singleton,
insert, insertList,
pop, popUntil, popWhile
) where
import Prelude
import qualified Data.Map as Map
import Data.Maybe
newtype Heap k v = Heap (Map.Map k [(k,v)])
empty :: Heap k v
empty = Heap Map.empty
fromList :: Ord k => [(k,v)] -> Heap k v
fromList xs = insertList xs empty
toList :: Heap k v -> [(k,v)]
toList (Heap mp) = concatMap reverse $ Map.elems mp
elems :: Heap k v -> [v]
elems (Heap mp) = concatMap (reverse . map snd) $ Map.elems mp
singleton :: Ord k => k -> v -> Heap k v
singleton k v = insert k v empty
insert :: Ord k => k -> v -> Heap k v -> Heap k v
insert k v (Heap xs) = Heap $ Map.insertWith (++) k [(k,v)] xs
insertList :: Ord k => [(k,v)] -> Heap k v -> Heap k v
insertList xs mp = foldr (uncurry insert) mp xs
pop :: Ord k => Heap k v -> Maybe ((k,v), Heap k v)
pop (Heap mp) | Map.null mp = Nothing
| null kvs = Just ((k1,v1), Heap mp2)
| otherwise = Just ((k1,v1), Heap $ Map.insert k kvs mp2)
where
((k,(k1,v1):kvs),mp2) = Map.deleteFindMin mp
popUntil :: Ord k => k -> Heap k v -> ([v], Heap k v)
popUntil x = popBy (< x)
popWhile :: Ord k => k -> Heap k v -> ([v], Heap k v)
popWhile x = popBy (<= x)
popBy :: Ord k => (k -> Bool) -> Heap k v -> ([v], Heap k v)
popBy cmp (Heap mp)
| Map.null mp || not (cmp k) = ([], Heap mp)
| otherwise = (reverse (map snd kvs) ++ res, mp3)
where
((k,kvs),mp2) = Map.deleteFindMin mp
(res,mp3) = popBy cmp (Heap mp2)