Discrete probability monad and an example analysis of the Monty Hall problem. Run it with `scala-cli monty_hall.sc`

monty_hall.sc

//> using dep "org.typelevel::spire:0.18.0"

import spire.math._
import spire.implicits._

case class Distro[A](entries: Map[A, Rational]):
  require(entries.values.qsum == Rational(1), s"Not normalized: $entries")

  def apply(value: A): Rational = entries.getOrElse(value, Rational(0))

  def map[B](f: A => B): Distro[B] = Distro(
    entries
      .groupMapReduce(entry => f(entry._1))(_._2)(_ + _)
      .toMap
  )

  def flatMap[B](f: A => Distro[B]): Distro[B] =
    val cases = 
      for (value, prob)   <- entries.toList
          (value2, prob2) <- f(value).entries
      yield (value2, prob * prob2)
    val aggregatedCases = cases.groupMapReduce(_._1)(_._2)(_ + _).toMap
    Distro(aggregatedCases)

  def withFilter(p: A => Boolean): Distro[A] =
    val filteredEntries = entries.filterKeys(p)
    val probMass = filteredEntries.values.qsum
    require(probMass > 0d)
    Distro(filteredEntries.map((value, prob) => value -> (prob/probMass)).toMap)

object Distro:
  def pure[A](value: A): Distro[A] = Distro(Map(value -> 1d))

  def uniform[A](values: Set[A]): Distro[A] = 	
    val prob = Rational(1, values.size)
    Distro(values.map(value => value -> prob).toMap)


val doors = Set(1, 2, 3)

val montyHallDistro =
  for carDoor <- Distro.uniform(doors)
      chosenDoor <- Distro.uniform(doors)
  yield (carDoor, chosenDoor)


println("Probability of winning")
val notSwitching = montyHallDistro.map((carDoor, chosenDoor) => carDoor == chosenDoor)
println(s"when not switching: ${notSwitching(true)}")

val switching =
  for (carDoor, chosenDoor) <- montyHallDistro
      revealedDoor <- Distro.uniform(doors - carDoor - chosenDoor)
      finalDoor = (doors - chosenDoor - revealedDoor).head
  yield carDoor == finalDoor
println(s"when switching: ${switching(true)}")