Discrete random variables in Haskell

pmf.hs

import Data.Either
import System.Random

-- |An association list mapping random variable outputs to their
-- probabilities, which should add up to 1.
type ProbMass a = [(a, Double)]

-- |Given a uniform random number in [0, 1) and a probability mass function
-- (PMF) describing the probabilities of various outcomes of a random
-- variable X, give a random outcome.
uniformToDiscrete :: Double -> ProbMass a -> a

-- This implementation is based on fitness-proportionate selection in
-- genetic algorithms.
--
-- Use a left fold to walk over the list of associations. The state that we
-- accumulate is Either a Double. We start out in (Right x) where the
-- starting value for x is the uniform random variable. We then decrease it
-- by each probability we encounter until it becomes negative. At this
-- point, we use (Left y) to indicate the return value.
--
-- Conceptually, imagine a roulette wheel where each number occupies
-- a space proportional to the likelihood of it being chosen. Unroll the
-- wheel and you get pockets spaced like:
--
-- 0.0       0.75       1.0
-- v           v         v
-- | 1 |  2   |  3   | 4 |
--
-- When we roll a number like 0.75, we skip as many elements from the front
-- as we can without the total probability of these pockets exceeding 0.75.
-- This is what the combining operation is doing: keeping track of a total
-- and then returning an outcome.
uniformToDiscrete rand pmf = fromLeft $ foldl combine (Right rand) pmf
  where
    combine :: Either a Double -> (a, Double) -> Either a Double
    combine (Right counter) (x, prob) = if (counter - prob) < 0
                                          then Left x
                                          else Right (counter - prob)
    combine acc@(Left _) _ = acc

    -- |A simple helper function which unwraps a Left
    fromLeft :: Either a b -> a
    fromLeft (Left x) = x


-- |Example which prints random strings forever.
example :: (RandomGen g) => g -> ProbMass String -> IO a
example gen pmf = let (uniform, gen') = random gen
                      x = uniformToDiscrete uniform pmf
                    in putStrLn x >> example gen' pmf

main = do
    gen <- getStdGen
    let pmf = [("x", 1.0/3.0),
               ("y", 2.0/3.0)]
    example gen pmf