haskell-unordered-containers / unordered-containers

Efficient hashing-based container types
BSD 3-Clause "New" or "Revised" License
221 stars 99 forks source link

More Data.Map compatibility: Add insertLookupWithKey #467

Open axman6 opened 2 years ago

axman6 commented 2 years ago

Following on from #172, adding insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k -> (Maybe a, Map k a) would be quite useful for avoiding multiple traversals of a map in various algorithms. An example I was working on recently was implementing the Misra-Gries summary for finding the most frequent elements in a streaming dataset:

misraGries :: Ord k => Int -> Fold k (M.Map k Int)
misraGries n = Fold step (M.empty :: M.Map k Int) id where
    step !old k = case M.insertLookupWithKey (\_k n _ -> n+1) k 1 old of
        (Just _prev, new) -> new
        (Nothing,    new)
            | M.size new < n -> new
            | otherwise -> M.mapMaybe (\case 1 -> Nothing; n -> Just (n-1)) old

I'd also like to have a HashMap based implementation but doing so involves several more traversals of the HashMap.

sjakobi commented 2 years ago

https://github.com/haskell-unordered-containers/unordered-containers/issues/245 is a feature request for a very similar function.

An implementation in terms of alterF should look similar to the one in https://github.com/haskell-unordered-containers/unordered-containers/issues/245#issuecomment-637399985.

Note that the current alterF implementation still takes two map traversals. A PR that optimizes the implementation would be very welcome though.

Since the implementation in terms of alterF is so simple, I'd currently prefer not to add a dedicated insertLookupWithKey function to the API.

axman6 commented 2 years ago

Interesting, I'll have to see if I can make it work with alterF