The "No Monty" problem (posed by Bartosz Milweski)

bartosz-nerd-snipe.hs

#!/usr/bin/env cabal

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

{- cabal:
build-depends:
    base
  , containers
  , mtl
  , random
-}

-- See: https://twitter.com/BartoszMilewski/status/1353804924472053760
--
-- A good chance to use some of Haskell's new random library!

-- Results:
--
-- Monty Hall Simulation
-- ---------------------
--
-- nTrials = 1000000
--
-- Original Monty Hall:
--      Keeping Door: P(Win) = 0.333089
--   Switching Doors: P(Win) = 0.666911
--
-- "No Monty" (including initial failures):
--      Keeping Door: P(Win) = 0.333305
--   Switching Doors: P(Win) = 0.332821
--
-- "No Monty" (excluding initial failures):
--      Keeping Door: P(Win) = 0.50049
--   Switching Doors: P(Win) = 0.49951

import Control.Monad (replicateM)
import Control.Monad.State.Strict (State)
import Data.Foldable (foldl')
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Word (Word8)
import System.Random (StdGen, mkStdGen)
import System.Random.Stateful
  ( RandomGen,
    RandomGenM,
    StateGen (StateGen),
    StateGenM,
    StatefulGen,
    Uniform,
    runStateGen_,
    uniformM,
    uniformRM,
  )

main :: IO ()
main = do
  let nTrials :: Int
      nTrials = 1_000_000

      g :: StdGen
      g = mkStdGen 42 -- 42 is just a seed
  putStrLn "Monty Hall Simulation"
  putStrLn "---------------------"
  putStrLn ""
  putStrLn $ "nTrials = " <> show nTrials
  putStrLn ""

  -- Original Monty Hall problem
  let originalKeepingDoorPWin :: PWin =
        runStateNTrials g nTrials KeepOriginalDoor originalMontyHallTrial
  let originalSwitchingDoorPWin :: PWin =
        runStateNTrials g nTrials SwitchDoors originalMontyHallTrial
  putStrLn "Original Monty Hall:"
  putStrLn $
    "     Keeping Door: P(Win) = "
      <> (show . unPWin $ originalKeepingDoorPWin)
  putStrLn $
    "  Switching Doors: P(Win) = "
      <> (show . unPWin $ originalSwitchingDoorPWin)
  putStrLn ""

  -- "No Monty" problem (including initial switch failure)
  let noMontyKeepingDoorPWin :: PWin =
        runStateNTrials g nTrials KeepOriginalDoor noMontyTrial
  let noMontySwitchingDoorsPWin :: PWin =
        runStateNTrials g nTrials SwitchDoors noMontyTrial
  putStrLn "\"No Monty\" (including initial failures):"
  putStrLn $
    "     Keeping Door: P(Win) = "
      <> (show . unPWin $ noMontyKeepingDoorPWin)
  putStrLn $
    "  Switching Doors: P(Win) = "
      <> (show . unPWin $ noMontySwitchingDoorsPWin)
  putStrLn ""

  -- "No Monty" problem (excluding initial switch failure)
  let noMontyNoFailKeepingDoorPWin :: PWin =
        runStateNTrials g nTrials KeepOriginalDoor noMontyNoFailTrial
  let noMontyNoFailSwitchingDoorsPWin :: PWin =
        runStateNTrials g nTrials SwitchDoors noMontyNoFailTrial
  putStrLn "\"No Monty\" (excluding initial failures):"
  putStrLn $
    "     Keeping Door: P(Win) = "
      <> (show . unPWin $ noMontyNoFailKeepingDoorPWin)
  putStrLn $
    "  Switching Doors: P(Win) = "
      <> (show . unPWin $ noMontyNoFailSwitchingDoorsPWin)

data Strategy = SwitchDoors | KeepOriginalDoor
  deriving (Eq, Show)

data Outcome = Win | NoWin
  deriving (Eq, Show)

data Door = Door1 | Door2 | Door3
  deriving (Eq, Ord, Show)

-- | Simulated approximation of probability of winning.
newtype PWin = PWin {unPWin :: Double}
  deriving (Eq, Ord, Show)

-- `Uniform` instance for random generation of `Door`s.
instance Uniform Door where
  uniformM g = word8ToDoor <$> uniformRM (1 :: Word8, 3 :: Word8) g
    where
      word8ToDoor :: Word8 -> Door
      word8ToDoor w = case w of
        1 -> Door1
        2 -> Door2
        3 -> Door3
        _ -> error "No values outside (1,3) should be generated!"

-- | A set containing all the doors.
allDoors :: Set Door
allDoors = Set.fromList [Door1, Door2, Door3]

-- | Pick a different door than the one passed-in.
pickOtherDoor :: (StatefulGen g m) => g -> Door -> m Door
pickOtherDoor g avoidDoor = do
  door :: Door <- uniformM g
  if door == avoidDoor
    then pickOtherDoor g avoidDoor
    else pure door

-- | Pick a random element from a non-empty set.
--
-- If the set is empty, a runtime error will occur.
--
-- Bit slow because it's indexing over lists. But they're tiny lists in this
-- use case.
randElem :: forall g m a. (StatefulGen g m) => g -> Set a -> m a
randElem _ s | Set.null s = pure $ error "randElem called on an empty set!"
randElem g s = do
  let list :: [a] = Set.toList s
  index <- uniformRM (0 :: Int, Set.size s - 1) g
  pure $ list !! index

-- | Run a single trial of the "original" Monty Hall problem, where
--   Monty ensures that the revealed door does NOT contain a win.
originalMontyHallTrial :: (StatefulGen g m) => g -> Strategy -> m Outcome
originalMontyHallTrial g strategy = do
  -- randomly assign a door which DOES contain the prize
  prizeDoor :: Door <- uniformM g

  -- randomly pick the contestant's chosen door
  initDoor :: Door <- uniformM g

  -- of the two other doors that the contestant didn't choose, Monty
  --  reveals one non-winning door.
  let remainingNonWinners :: Set Door
      remainingNonWinners =
        allDoors `Set.difference` (Set.fromList [prizeDoor, initDoor])
  montyRevealDoor :: Door <- randElem g remainingNonWinners

  -- now we finalize the contestant's door. If they are using the
  --  SwitchDoors strategy then we switch to the remaining door.
  --  Otherwise, we keep their original door.
  let finalDoor :: Door
      finalDoor = case strategy of
        KeepOriginalDoor -> initDoor
        SwitchDoors ->
          head . Set.toList $
            allDoors `Set.difference` (Set.fromList [initDoor, montyRevealDoor])

  -- finally evaluate how we did. If the finalDoor matches the
  -- prizeDoor then the contestant won, otherwise they didn't.
  pure $ if finalDoor == prizeDoor then Win else NoWin

-- | Run a single trial of the "no Monty" problem, as described by
--   Bartosz Milewski.
--
-- In this version, we quit with a failure if we're "out of the game".
noMontyTrial :: (StatefulGen g m) => g -> Strategy -> m Outcome
noMontyTrial g strategy = do
  -- randomly assign a door which DOES contain the prize
  prizeDoor :: Door <- uniformM g

  -- randomly pick the contestant's chosen door
  initDoor :: Door <- uniformM g

  -- randomly pick the contestant's second door
  let remainingDoors :: Set Door
      remainingDoors = allDoors `Set.difference` (Set.fromList [initDoor])
  secondDoor :: Door <- randElem g remainingDoors

  -- If the second door had the car, then we're out of the game, and we
  -- lost at this point.
  if secondDoor == prizeDoor
    then pure NoWin
    else do
      -- otherwise we're still in the game
      let finalDoor :: Door
          finalDoor = case strategy of
            KeepOriginalDoor -> initDoor
            SwitchDoors ->
              head . Set.toList $
                allDoors `Set.difference` (Set.fromList [initDoor, secondDoor])

      pure $ if finalDoor == prizeDoor then Win else NoWin

-- | Run a single trial of the "no Monty" problem, as described by
--   Bartosz Milewski.
--
-- In this version, we don't include failures due to being wrong about the
-- second door pick. The "if you're still in the game" part.
noMontyNoFailTrial :: (StatefulGen g m) => g -> Strategy -> m Outcome
noMontyNoFailTrial g strategy = do
  -- randomly assign a door which DOES contain the prize
  prizeDoor :: Door <- uniformM g

  -- randomly pick the contestant's chosen door
  initDoor :: Door <- uniformM g

  -- randomly pick the contestant's second door
  let remainingDoors :: Set Door
      remainingDoors = allDoors `Set.difference` (Set.fromList [initDoor])
  secondDoor :: Door <- randElem g remainingDoors

  -- If the second door had the car, then we're out of the game, and in this
  -- version, we'll just try again until we get a trial where that didn't
  -- happen.
  if secondDoor == prizeDoor
    then noMontyNoFailTrial g strategy
    else do
      -- otherwise we're still in the game
      let finalDoor :: Door
          finalDoor = case strategy of
            KeepOriginalDoor -> initDoor
            SwitchDoors ->
              head . Set.toList $
                allDoors `Set.difference` (Set.fromList [initDoor, secondDoor])

      pure $ if finalDoor == prizeDoor then Win else NoWin

runNTrials ::
  (StatefulGen g m) =>
  g ->
  Int ->
  Strategy ->
  (g -> Strategy -> m Outcome) ->
  m PWin
runNTrials g n strategy trial = do
  trials :: [Outcome] <- replicateM n (trial g strategy)
  let nDouble :: Double = fromIntegral n
  let isWin :: Outcome -> Bool
      isWin o = o == Win
  pure $ PWin $ fromIntegral (count isWin trials) / nDouble

runStateNTrials ::
  ( RandomGen g,
    h ~ StateGenM (StateGen g),
    m ~ State (StateGen g)
  ) =>
  g ->
  Int ->
  Strategy ->
  (h -> Strategy -> m Outcome) ->
  PWin
runStateNTrials g nTrials strategy trial =
  runStateGen_ (StateGen g) (\g' -> runNTrials g' nTrials strategy trial)

-- | Count items in a list which satisfy a predicate.
count :: forall a. (a -> Bool) -> [a] -> Int
count p xs = foldl' f 0 xs
  where
    f :: Int -> a -> Int
    f n x = if p x then n + 1 else n