Skip to content

Instantly share code, notes, and snippets.

@sortega
Created February 22, 2023 17:05
Show Gist options
  • Save sortega/00fddef438c096c3ece718c3a231757b to your computer and use it in GitHub Desktop.
Save sortega/00fddef438c096c3ece718c3a231757b to your computer and use it in GitHub Desktop.
Discrete probability monad and an example analysis of the Monty Hall problem. Run it with `scala-cli 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)}")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment