Red-Green vs Red-Red

Main.hs

module Main where

{-
This is a simulation of a problem posed by Daniel Litt in a tweet:
https://twitter.com/littmath/status/1751648838501224790

> You are given an urn containing 100 balls; n of them are red, and
> 100-n are green, where n is chosen uniformly at random in [0, 100].
> You take a random ball out of the urn—it’s red—and discard it. The
> next ball you pick (out of the 99 remaining) is:
> More likely to be red | More likely to be green | Equal | Results

My answer was "more likely to be red." The simulation below verifies
this experimentally, with an apparent 2:1 likelihood in favor of red.
-}

import Control.Monad (replicateM)
import Data.MultiSet
    ( delete, empty, insertMany, occur, size, MultiSet )
import System.Random (randomRIO)

data Ball = Red | Green deriving (Eq, Ord)
type Urn = MultiSet Ball
data Outcome = GreenHalt | RedGreen | RedRed deriving Eq

-- Make an urn with r Red balls and 100-r Green balls
makeUrn :: Int -> Urn
makeUrn r = insertMany Red r $ insertMany Green (100 - r) empty

-- Make a random urn with [0,100] red balls, equal likelihood
makeUrnRand :: IO Urn
makeUrnRand = makeUrn <$> randomRIO (0, 100)

-- Remove a random ball from an urn, equal likelihood
removeBallRand :: Urn -> IO (Ball, Urn)
removeBallRand urn = do
    let r = occur Red urn -- how many reds?
    i <- randomRIO (0, size urn - 1) -- random selection from urn
    pure $ if i < r -- if we chose red in this round
        then (Red, delete Red urn) -- remove one red ball
        else (Green, delete Green urn) -- remove one green ball

-- Simulate drawing a ball (and again if the first was red)
simulate :: IO Outcome
simulate = do
    urn <- makeUrnRand
    (ball1, smallerUrn) <- removeBallRand urn
    case ball1 of
      Green -> pure GreenHalt -- if ball drawn is green, halt sim
      Red -> do -- first ball drawn is red, now draw another
        (ball2, _) <- removeBallRand smallerUrn
        case ball2 of
          Red -> pure RedRed -- second ball drawn was red
          Green -> pure RedGreen -- second ball drawn was green

-- Simulate N times
simulateMany :: Int -> IO [Outcome]
simulateMany n = replicateM n simulate

-- Generate report from multiple sims
summarize :: [Outcome] -> String
summarize outcomes =
    let redReds = length $ filter (== RedRed) outcomes
        redGreens = length $ filter (== RedGreen) outcomes
    in  "There were " <>
        show redReds <>
        " red-reds and " <>
        show redGreens <>
        " red-greens."

main :: IO ()
main = do
    outcomes <- simulateMany 9000
    print $ summarize outcomes