mystery.scala

// 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