CODEWITHSUNDEEP
haskellfold
sara
sara

Reputation: 3589

Haskell - Sum up probability list

I'm going through Learn You A Haskell and just finished the part "For a few monads more". In this part we created a newtype Prob a = Prob { getProb :: [(a, Rational)] } and created a Monad instance for it. This allows us to calculate probabilities of results in non-deterministic calculations like the following:

data Coin = Heads | Tails deriving (Show, Eq)

coin :: Prob Coin
coin = Prob [(Heads, 1%2), (Tails, 1%2)]

loadedCoin :: Prob Coin
loadedCoin = Prob [(Heads, 1%10), (Tails, 9%10)]

coinTest :: Prob Bool
coinTest = do
    a <- coin
    b <- coin
    c <- loadedCoin
    return (all (==Tails) [a,b,c])

Of course, this doesn't generate a very pretty result:

getProb coinTest
>> [(False,1 % 40),(False,9 % 40),(False,1 % 40),(False,9 % 40),(False,1 % 40),(False,9 % 40),(False,1 % 40),(True,9 % 40)]

It was left as an exercise to the reader to write a neat function to sum up all the Falses and all the Trues so we get [(True,9 % 40),(False,31 % 40)]. I've managed to do this, sort of. It works for this particular case, but I feel like it isn't a useful function at all since it's so specialized. Here's what I came up with:

sumProbs :: Prob Bool -> Prob Bool
sumProbs (Prob ps) = let (trues, falses) = partition fst ps
                         ptrue = reduce trues
                         pfalse = reduce falses
                     in Prob [ptrue, pfalse]
                     where reduce = foldr1 (\(b,r) (_,r') -> (b,r+r'))

I'd love to generalize it to work for any Eq a => Prob a, but so far no luck. I was thinking of maybe using a Map together with unionWith or something like that. Or maybe I can exploit the fact that (a,b) has a Functor b instance? I think I'm missing some simpler more elegant solution.

So, to sum it up: How can I write a function sumProbs :: (Eq a) => Prob a -> Prob a that sums up all the probabilities that share the same value (key)?

Upvotes: 2

Views: 235

Answers (2)

behzad.nouri
behzad.nouri

Reputation: 78001

If you use Data.Map, then fromListWith and toList will do:

import Data.Map (toList, fromListWith)

newtype Prob a = Prob { getProb :: [(a, Rational)] }
    deriving Show

sumProbs :: (Ord a) => Prob a -> Prob a
sumProbs = Prob . toList . fromListWith (+) . getProb

Relaxing Ord a to Eq a would require a less efficient quadratic computation; something like:

sumProbs :: (Eq a) => Prob a -> Prob a
sumProbs = Prob . foldr go [] . getProb
    where
    go (x, y) = run
        where
        run [] = (x, y):[]
        run ((a, b):rest)
            | x == a    = (x, y + b): rest
            | otherwise = (a, b): run rest

Upvotes: 3

leftaroundabout
leftaroundabout

Reputation: 120741

Using Map is a good idea, but you'll need Ord a in addition to Eq a. If you're ok with that, then we can also do simpler list solution: just replace partition with a combination of sortBy and groupBy:

import Data.List (groupBy, sortBy)
import Data.Function (on)

sumProbs :: (Ord a) => Prob a -> Prob a
sumProbs (Prob ps) = Prob . map reduce
                   . groupBy ((==)`on`fst)
                   $ sortBy (compare`on`fst) ps

Upvotes: 2

Related Questions