Added average and average1 functions

See kowainik#316.
aleator · Aug 12, 2020 · b33012a · b33012a
1 parent 48d2160
commit b33012a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/src/Relude/Extra/Foldable.hs b/src/Relude/Extra/Foldable.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE Safe #-}
+{-# LANGUAGE BangPatterns #-}
 
 {- |
 Copyright:  (c) 2018-2020 Kowainik
@@ -14,6 +15,7 @@ Contains utility functions for working with tuples.
 
 module Relude.Extra.Foldable
     ( foldlSC
+    , average
     ) where
 
 import Relude
@@ -38,3 +40,15 @@ foldlSC f = flip $ foldr go id
         Left l  -> l
         Right r -> k r
 {-# INLINE foldlSC #-}
+
+{- | Compute average of 'Foldable' 
+-}
+average :: (Foldable f, Fractional a) => f a -> Maybe a
+average xs
+    | null xs = Nothing
+    | otherwise = Just 
+                    . uncurry (/) 
+                    . foldl' (\(!total, !count) x -> (total + x, count + 1)) (0,0) 
+                    $ xs
+{-# INLINE average #-}
+
diff --git a/src/Relude/Extra/Foldable1.hs b/src/Relude/Extra/Foldable1.hs
@@ -5,6 +5,7 @@
 {-# LANGUAGE TypeFamilies         #-}
 {-# LANGUAGE TypeOperators        #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE BangPatterns         #-}
 
 {- |
 Copyright:  (c) 2011-2015 Edward Kmett
@@ -26,6 +27,7 @@ contradiction with 'Data.Foldable.Foldable'.
 module Relude.Extra.Foldable1
     ( Foldable1 (..)
     , foldl1'
+    , average1
     ) where
 
 import Relude hiding (Product (..), Sum (..))
@@ -370,3 +372,10 @@ foldl1' :: (a -> a -> a) -> NonEmpty a -> a
 foldl1' _ (x :| [])     = x
 foldl1' f (x :| (y:ys)) = foldl' f (f x y) ys
 {-# INLINE foldl1' #-}
+
+{- | Compute average of 'Foldable1' 
+-}
+average1 :: (Foldable1 f, Fractional a) => f a -> a
+average1 = uncurry (/) . foldl' (\(!total, !count) x -> (total + x, count + 1)) (0,0) 
+{-# INLINE average1 #-}
+