BUCKETS - Editorial





Author: AmirReza PoorAkhavan
Tester: Danial Erfanian
Editorialist: Michael Nematollahi


DP, Probabilities


We’re going to use Dynamic Programming. Define p*[j] = the probability that the ball we pick from the i^{th} bucket is of color j. If we can quickly fill this dp array, the answer will obviously be p[N][j] for each color j.


Let’s use Dynamic Programming. Define p*[j] as the probability that the ball we pick from the i^{th} bucket is of color j. Also, let’s define sm* as the total number of balls in the i^{th} bucket at the beginning. Note that at the i^{th} step, the number of balls in the i^{th} bucket is sm*+1 if i > 1 and sm* otherwise.
The above explanations lead to a solution to calculate p*[j] that is as follows:

for i=1, basic knowledge of the probability theory gives us:
p[1][j] = \frac{a[1][j]}{sm[1]}

for i>1, we have:
p*[j] = (\sum_{w=1, w!=j}^{K} p[i-1][w] imes \frac{a*[j]}{sm*+1}) + p[i-1][j] imes \frac{a*[j]+1}{sm*+1}

we can rewrite the above formula as:
p*[j] = (\frac{a*[j]}{sm*+1} imes \sum_{w=1, w!=j}^{K} p[i-1][w]) + p[i-1][j] imes \frac{a*[j]+1}{sm*+1}
which based on the additivity axiom is equal to:
p*[j] = (\frac{a*[j]}{sm*+1} imes (1 - p[i-1][j])) + p[i-1][j] imes \frac{a*[j]+1}{sm*+1}
which can easily be calculated in O(N imes K).

Refer to the editorialist’s code to see the described solution. Note that the code contains some small tricks to make the implementation easier.


Author’s solution can be found here

Tester’s solution can be found here

Editorialist’s solution can be found here.


I didn’t understand the question. Was he picking just one ball from bucket i and moving to i+1 OR was he picking all the balls from bucket i and then moving to i+1 ?


@dp1priyesh3_3, since I can’t post a comment on your post, I’ll add it here.

It essentially a step-wise process, in step i, you choose a ball from bucket i and then go ahead and place the ball in bucket i + 1. In step i + 1, you pick a ball from bucket i + 1 (this ball can be the one we added in step i too) and place it in bucket i + 2 and so on…

There’s a far simpler solution though I’m not sure if it counts as DP or not. It probably does, but can’t say.

So, essentially we have to compute the average of probability in all possible cases which is the same as calculating the probability in the average case.

So, what would be the average case?

Let’s label some stuff first! It always helps in solving problems.

Let B_i refer to bucket i, a_{i, j} be the number of balls of colour j in B_i and S_i = \sum_{j=1}^k a_{i, j} be the total number of balls in the B_i.

Well, so after step i - 1 what will the distribution of B_i look like? There will clearly be S_{i - 1} possible cases, since we are picking up a ball from B_{i - 1}.

In a_{i - 1, j} of the cases the value of a_{i, j} is increased by 1 and value of a_{i, r} (r eq j) remains the same.

So, the distribution of color j in the average case will simply be

a_{i, j} + \frac{a_{i - 1, j}}{S_i}.

Now, simply loop this from i = 2 to n. The pseudo-code is:

    for i = 1 to n do
        S* := 0
        for j = 1 to k do
            S* → S* + a*[j]
    for i = 2 to n do
        for j = 1 to k do
            a*[j] → a*[j] + a[i - 1][j]/S[i - 1]
        S* → S* + 1
    for j = 1 to k do
        print a[n][j]/S[n]

Also, notice the line S_i o S_i + 1, don’t forget to increment S_i by 1 because we just added a ball.

Here’s my implementation in C++ which is what I submitted during the contest.

Don’t forget that a_{i, j} has to be floating-point. :stuck_out_tongue:

Side-note: If you are looking for a rigorous proof of the average case shenanigans, search for Expected Value and Linearity of Expectation.
Such an obvious statement gives many unexpected results.

Happy New Year!


ohh… seems scary after looking at DP tag… though I solved it XD