Skip to content

Instantly share code, notes, and snippets.

@jliszka
Last active November 30, 2018 15:33
Show Gist options
  • Save jliszka/4ea988d9854e34aa1e1263f64f15fa9c to your computer and use it in GitHub Desktop.
Save jliszka/4ea988d9854e34aa1e1263f64f15fa9c to your computer and use it in GitHub Desktop.
// Multinomial formula
def multi(ks: Vector[Int]): Double = {
(1 to ks.sum)
.map(_ / ks.size.toDouble)
.zip(ks.flatMap(1 to _))
.map({ case (n, d) => n / d }).product
}
// Expected number of cereal boxes you have to buy if you want to collect at least `m` of each of `n` surprise mystery toys
def e(n: Int, m: Int, M: Int = 20, v: Vector[Int] = Vector.empty): Double = {
if (n == 1) multi(v :+ (m-1)) * (v.sum + m)
else (m to M).map(i => e(n-1, m, N, v :+ i)).sum
}
// Same thing but done by random sampling
def g(n: Int, m: Int): Distribution[Double] = {
discreteUniform(1 to n)
.until(xs => xs.groupBy(x => x).values.map(_.size).count(_ >= m) == n)
.map(_.size.toDouble)
}
// Special case of m=1
def f(n: Int) = (1 to n).map(i => 1.0 * n / i).sum
// e reproduces f when m=1
scala> f(3)
res0: Double = 5.5
scala> e(3, 1, 20)
res1: Double = 5.499903100004326
// e matches g
scala> g(5, 4).ev
res2: Double = 32.897099999999995
scala> e(5, 4, 20)
res2: Double = 32.44101314690209
scala> g(2, 10).ev
res3: Double = 23.536799999999964
scala> e(2, 10, 50)
res4: Double = 23.523939115154043
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment