module Statistics.Iteratee.Sample (
minMaxNBy
, range
, mean
, harmonicMean
, variance
, stdDev
) where
import Statistics.Iteratee.Compat
import Control.Arrow
import Control.Monad
import Data.Iteratee as I
import Data.Heap as Heap
minMaxNBy
:: forall m prio s el. (Monad m, Ord prio, ListLikey s el)
=> Int
-> (el -> prio)
-> Iteratee s m ([(prio,el)],[(prio,el)])
minMaxNBy ns prio = finalize `liftM` I.foldl' step (Heap.empty,Heap.empty)
where
finalize :: (MaxPrioHeap prio el, MinPrioHeap prio el)
-> ([(prio,el)],[(prio,el)])
finalize = Heap.toDescList *** Heap.toDescList
addHeap val = Heap.insert (prio val, val)
step :: (MaxPrioHeap prio el, MinPrioHeap prio el) -> el
-> (MaxPrioHeap prio el, MinPrioHeap prio el)
step (!mins,!maxes) val = let sz = Heap.size mins
adj hp = if sz >= ns
then Heap.drop 1 hp
else hp
in (adj $ addHeap val mins
,adj $ addHeap val maxes)
range :: (Monad m, ListLikey s el, Num el, Ord el)
=> Iteratee s m el
range = finalize `liftM` minMaxNBy 1 id
where
finalize ([mins],[maxes]) = snd maxes snd mins
finalize _ = 0
mean :: forall s m el. (Fractional el, Monad m, ListLikey s el)
=> Iteratee s m el
mean = fst `liftM` I.foldl' step (0,0)
where
step :: (el,Integer) -> el -> (el,Integer)
step (!m,!n) x = let m' = m + (xm) / fromIntegral n'
n' = n + 1
in (m',n')
harmonicMean :: (Fractional el, Monad m, ListLikey s el) => Iteratee s m el
harmonicMean = finalize `liftM` I.foldl' step (0,0 :: Integer)
where
finalize (m,n) = fromIntegral n / m
step (!m,!n) val = (m+(1/val),n+1)
var :: (Fractional el, Integral t, Monad m, ListLikey s el)
=> Iteratee s m (t, el, el)
var = I.foldl' step (0,0,0)
where
step (!n,!m,!s) x = let n' = n+1
m' = m+d/fromIntegral n'
s' = s+d* (xm')
d = xm
in (n',m',s')
variance :: (Fractional b, Monad m, ListLikey s b) => Iteratee s m b
variance = finalize `liftM` var
where
finalize (n,_,s)
| n > 1 = s / fromInteger n
| otherwise = 0
stdDev :: (Floating b, Monad m, Functor m, ListLikey s b) => Iteratee s m b
stdDev = sqrt `liftM` variance