← All solutions

Shortest judge solution: 1278 bytes. Shortest team solution (during contest): 1050 bytes.

Let us number the islanders 1, 2, . . . , n in some fixed order. Let (g₁, g₂, . . . , g_n) be any integers summing up to d + n, with g_i ≥ 1 – a potential distribution of gems after d nights. It turns out the probability the gems are distributed that way after d nights is the same for any distribution.

Let us prove this by induction on d. For d = 0, there is only one distribution (1, 1, . . . , 1). For a higher d, there are n potential ways to get to (g₁, g₂, . . . , g_n) – from any distribution where exactly one islander has exactly one gem less. The probability that we get from (g₁, g₂, . . . , g_k − 1, . . . , g_n) to (g₁, g₂, . . . , g_n) is (g_k − 1)/(n + d − 1) – one of the g_k − 1 has to be the one to split.

Since the probability for each distribution for d − 1 is some fixed p by the inductive hypothesis, the total probability of getting to (g₁, g₂, . . . , g_n) is p · ∑_k (g_k − 1)/(n + d − 1) = p · (n + d − 1)/(n + d − 1) = p.

And the sum of all g_k − 1 is d – the number of gems over the first one that all the islanders have (technically, we also need to exclude the distributions where g_k = 1 from the sum, since (g₁, g₂, . . . , 0, . . . , g_n) is not a valid distribution, but the probability of that distribution is multiplied by g_k − 1 = 0 in the sum anyway, so it does not matter).

With that observation, we are prepared to calculate the expected value we are being asked for. Since all the distributions are equally probable, we can instead calculate the sum of the number of gems the r richest islanders have, over all possible distributions, and then divide by the number of all possible distributions. The number of ways to assign d gems to n people is ^n+d−1C_d, and we can precalculate all the binomial coefficients in quadratic time and memory, so let’s just look at the sum.

Let us denote the sum over all legal distributions of n + d gems to n people of the number of gems the r richest people have by S(n, d). We can write a recursion for S(n, d). If n ≤ r, S(n, d) = n + d.

For larger n, there is a contribution coming from the first gems of every islander – that is obviously r · ^n+d−1C_d. On top of that, we have some islanders that have only one gem – their contribution is finished – and some that have more gems. These islanders that have more gems, after removing their first gem, still have a valid distribution of the remaining d gems. So, if we fix g – the number of islanders that have exactly one gem – the remaining contribution is ⁿC_g S(n − g, d − n + g). This gives us a recursion formula for S(n, d):

S(n, d) = ^n+d−1C_d ( r + ∑_g=0^{n n}C_g S(n − g, d − n + g) )

This recursion formula translates to a dynamic program for calculating S(n, d) in time O(n²d), which is fast enough. Even less efficient solutions (with an extra log factor in there) had a chance to pass if efficiently implemented.

This problem can be solved even faster. For instance, a quadratic-time solution, communicated to us by Petr Mitrichev, goes roughly as follows. Let V(k, y) be the number of distributions for which exactly k islanders have at least y gems. We have S(n, d) = ∑_{k,y ≥ 1} min(k, r) V(k, y).

If we denote by W(k, y) the number of distributions where at least k islanders have at least y gems, then V(k, y) = W(k, y) − W(k + 1, y). And W(k, y) can be calculated from the inclusion-exclusion principle as

W(k, y) = ∑_{ℓ ≥ k} (−1)^{ℓ n}C_k−1 ^ℓ−kC_ℓ−1 ^{n+d−ℓ(y−1)−1}C_{d−ℓ(y−1)}

The proofs of these equalities, as well as of the fact that the resulting algorithm is quadratic, are left as an exercise to the reader.