Tester: Sergey Kulik
Editorialist: Lalit Kundu
Combinatorics, Dynamic Programming, Probability
Given N ballons, i’th balloon has color Ci and cost Pi dollars. We choose any subset (chosen randomly, possibly empty) of the balloons such that the number of different colors in that subset is at least M. What is the expected cost of such a subset?
Since subsets are randomly chosen, expected cost is (total cost of subsets with at least M colors)/(total number of subsets with atleast M colors).
Let’s keep colors[i]=number of balloons with color=i and cost[i]=sum of cost of all balloons with color=i.
Let color=number of distinct colors.
Let solve the denominator first.
We keep DP1[i][j], which denotes number of subsets which consists of colors 1 to i, and have j distinct colors.
DP1[i][j]= // we are not using the i'th color DP1[i-1][j] + // we use the i'th color // now, we can choose any combination of colors[i] balloons. // but not the one in which no balloon is present. // therefore, we have (2**colors[i])-1 different ways. DP1[i-1][j-1]*((2**colors[i])-1)
Now, we come to the numerator.
For numerator, we keep DP2[i][j], which denotes the total cost of all subsets which consist of colors 1 to i, and have j distinct colors.
DP2[i][j]= // we are not using the i'th color DP2[i-1][j] + // we use the i'th color // firstly, we again have (2**colors[i])-1 options // so we add DP2[i-1][j-1]*((2**colors[i])-1) + //also, we add //total sum for selecting i'th color * number of ways of selecting j-1 colors from i-1 colors(DP1[i-1][j-1]) // total sum for selecting i'th color is cost[i]*(2**(colors[i]-1)). How? DP1[i-1][j-1]*cost[i]*(2**(colors[i]-1))
But, we count only those subsets where different colors is more than or equal to M.
num=den=0 for i=M to color: num += DP2[color][i] den += DP1[color][i] print num/den
Complexity: O(color*color) where color=number of distinct colors.