minor additions to PQueue

This commit is contained in:
Jack Wines 2024-04-15 02:27:14 -07:00
parent 37e28f8c0e
commit 110c6eb8a5
No known key found for this signature in database
GPG key ID: 25B20640600571E6

View file

@ -8,7 +8,7 @@ import Data.List.NonEmpty qualified as NE
newtype PQueue a = PQueue newtype PQueue a = PQueue
{ toMap :: IM.IntMap (NE.NonEmpty a) { toMap :: IM.IntMap (NE.NonEmpty a)
} }
deriving (Functor, Show) deriving (Functor, Show, Eq, Ord)
fromList :: [(Int, a)] -> PQueue a fromList :: [(Int, a)] -> PQueue a
fromList = PQueue . IM.fromListWith NE.append . map (Bi.second NE.singleton) fromList = PQueue . IM.fromListWith NE.append . map (Bi.second NE.singleton)
@ -16,11 +16,23 @@ fromList = PQueue . IM.fromListWith NE.append . map (Bi.second NE.singleton)
singleton :: Int -> a -> PQueue a singleton :: Int -> a -> PQueue a
singleton key a = PQueue $ IM.singleton key [a] singleton key a = PQueue $ IM.singleton key [a]
minView :: PQueue a -> Maybe (Int, (a, PQueue a)) abstractView :: (IM.IntMap (NE.NonEmpty a) -> Maybe (IM.Key, NE.NonEmpty a)) -> PQueue a -> Maybe (IM.Key, (a, PQueue a))
minView (PQueue m) = case IM.lookupMin m of abstractView f (PQueue m) = case f m of
Nothing -> Nothing Nothing -> Nothing
(Just (key, x NE.:| (x' : xs))) -> Just (key, (x, PQueue $ IM.insert key (x' NE.:| xs) m)) (Just (key, x NE.:| (x' : xs))) -> Just (key, (x, PQueue $ IM.insert key (x' NE.:| xs) m))
(Just (key, x NE.:| [])) -> Just (key, (x, PQueue $ IM.delete key m)) (Just (key, x NE.:| [])) -> Just (key, (x, PQueue $ IM.delete key m))
minView :: PQueue a -> Maybe (Int, (a, PQueue a))
minView = abstractView IM.lookupMin
maxView :: PQueue a -> Maybe (Int, (a, PQueue a))
maxView = abstractView IM.lookupMax
insert :: Int -> p -> PQueue p -> PQueue p insert :: Int -> p -> PQueue p -> PQueue p
insert key a = PQueue . IM.insertWith NE.append key (NE.singleton a) . toMap insert key a = PQueue . IM.insertWith NE.append key (NE.singleton a) . toMap
elems :: PQueue a -> [a]
elems = concatMap NE.toList . IM.elems . toMap
keys :: PQueue a -> [IM.Key]
keys = IM.keys . toMap