@@ -35,7 +35,6 @@ module Streamly.Internal.Data.Fold.Window
3535 --
3636 lmap
3737 , cumulative
38− , nub
3938
4039 -- ** Sums
4140 , length
@@ -49,6 +48,9 @@ module Streamly.Internal.Data.Fold.Window
4948 , maximum
5049 , range
5150 , mean
51+
52+ -- ** Distribution
53+ , frequency
5254 )
5355where
5456
@@ -66,10 +68,12 @@ import qualified Streamly.Internal.Data.Ring.Foreign as Ring
6668
6769import Prelude hiding (length, sum, minimum, maximum)
6870
69−
7071-- $setup
7172-- >>> import Data.Bifunctor(bimap)
7273-- >>> import qualified Streamly.Data.Fold as Fold
74+-- >>> import qualified Streamly.Internal.Data.Fold.Window as FoldW
75+-- >>> import qualified Streamly.Internal.Data.Ring.Foreign as Ring
76+-- >>> import qualified Streamly.Prelude as Stream
7377-- >>> import Prelude hiding (length, sum, minimum, maximum)
7478
7579-------------------------------------------------------------------------------
@@ -313,38 +317,35 @@ maximum n = fmap (fmap snd) $ range n
313317mean :: forall m a. (Monad m, Fractional a) => Fold m (a, Maybe a) a
314318mean = Fold.teeWith (/) sum length
315319
316−-- | Drop repeated elements within the specified
317−-- tumbling window in the stream.
318−-- /Pre-release/
320+-------------------------------------------------------------------------------
321+-- Distribution
322+-------------------------------------------------------------------------------
323+
324+-- XXX We can use a Windowed classifyWith operation, that will allow us to
325+-- express windowed frequency, mode, histograms etc idiomatically.
326+
327+-- | Count the frequency of elements in a sliding window.
319328--
320−-- >>> import qualified Streamly.Prelude as S
321−-- >>> S.fold nub $ S.fromList [(1, Nothing), (1, Nothing), (3, Nothing), (4, Just 1)]
322−-- fromList [(1,1),(3,1),(4,1)]
329+-- >>> input = Stream.fromList [1,1,3,4,4::Int]
330+-- >>> f = Ring.slidingWindow 4 FoldW.frequency
331+-- >>> Stream.fold f input
332+-- fromList [(1,1),(3,1),(4,2)]
323333--
324−{-# INLINE nub #-}
325−nub :: (Monad m, Ord a) => Fold m (a, Maybe a) (Map.Map a Int)
326−nub = Fold step initial extract
334+-- /Pre-release/
335+--
336+{-# INLINE frequency #-}
337+frequency :: (Monad m, Ord a) => Fold m (a, Maybe a) (Map.Map a Int)
338+frequency = Fold.foldl' step Map.empty
327339
328340 where
329341
330− initial =
331− return $ Partial Map.empty
332−
333− updateRefCounts new map0 = return $ Map.insertWith (+) new (1::Int) map0
334−
335− alter a m1 (Just v) =
342+ decrement v =
336343 if v == 1
337− then do
338− _ <- return $ Map.delete a m1
339− return Nothing
340− else return $ Just (v -1)
341− alter _ _ Nothing = error "Invalid old element"
342−
343− step map0 (new, mOld) = do
344− m1 <- updateRefCounts new map0
345− Partial <$>
346− case mOld of
347− Just a -> Map.alterF (alter a m1) a m1
348− Nothing -> return m1
349−
350− extract = return
344+ then Nothing
345+ else Just (v - 1)
346+
347+ step refCountMap (new, mOld) =
348+ let m1 = Map.insertWith (+) new 1 refCountMap
349+ in case mOld of
350+ Just k -> Map.update decrement k m1
351+ Nothing -> m1