PROBLEM LINK:
Author: AmirReza PoorAkhavan
Tester: Danial Erfanian
Editorialist: Michael Nematollahi
PREREQUISITES:
DP, Probabilities
QUICK EXPLANATION:
We’re going to use Dynamic Programming. Define p[i][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.
EXPLANATION:
Let’s use Dynamic Programming. Define p[i][j] as the probability that the ball we pick from the i^{th} bucket is of color j. Also, let’s define sm[i] 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[i]+1 if i > 1 and sm[i] otherwise.
The above explanations lead to a solution to calculate p[i][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[i][j] = (\sum_{w=1, w!=j}^{K} p[i-1][w] \times \frac{a[i][j]}{sm[i]+1}) + p[i-1][j] \times \frac{a[i][j]+1}{sm[i]+1}
we can rewrite the above formula as:
p[i][j] = (\frac{a[i][j]}{sm[i]+1} \times \sum_{w=1, w!=j}^{K} p[i-1][w]) + p[i-1][j] \times \frac{a[i][j]+1}{sm[i]+1}
which based on the additivity axiom is equal to:
p[i][j] = (\frac{a[i][j]}{sm[i]+1} \times (1 - p[i-1][j])) + p[i-1][j] \times \frac{a[i][j]+1}{sm[i]+1}
which can easily be calculated in O(N \times 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 AND TESTER’S SOLUTIONS:
Author’s solution can be found here
Tester’s solution can be found here
Editorialist’s solution can be found here.