Dwango Programming Contest 6th Task C. Cookie Distribution

The outcome of an experiment can be represented as N sets D[1], ..., D[N], with each set D[i] being a subset of {1, 2, ..., K} denoting days on which the i-th child was chosen to get a cookie. We have c[i] = |D[i]| for i = 1, 2, ..., N.

We only need to consider outcomes where each child was chosen at least once during the K days, in ohter words, no set D[i] is empty. For an outcome, the quantity c[1] * c[2] * ... * c[N] can be interpreted as the number of N-tuples (d[1], d[2], ..., d[N]) where d[i] ∈ D[i] for i = 1, 2, ..., N.

What we need to compute is number of N-tuples (d[1], d[2], ..., d[N]) in all outcomes, where d[i] means on which day the i-th child got a cookie.

Now consider in how many outcomes is there a particular N-tuple (d[1], d[2], ..., d[N]) as described. The answer turns out to be C(N - b[1], a[1] - b[1]) * C(N - b[2], a[2] - b[2]) * ... * C(N - b[K], a[K] - b[K]) where a[i] means number of children chosen on the i-th day, and b[i] means the number of times that i appears in the N-tuple (d[1], d[2], ..., d[N]).

We see that the number that a particular N-tuple (d[1], d[2], ..., d[N]) is counted is determined by the corresponding K-tuple (b[1], b[2], ..., b[K]).

Now consider how many N-tuple (d[1], d[2], ..., d[N]) corresponds to a particular K-tuple (b[1], b[2], ..., b[K]). This problem can be rephrased as follows: distribute N children over K days, such that there are b[i] children on the i-th day, for i = 1, ..., K. This is a classical problem and the answer is N! / (b[1]! * b[2]! * ... * b[K]!).

Now we need to add up C(N - b[1], a[1] - b[1]) * C(N - b[2], a[2] - b[2]) * ... * C(N - b[K], a[K] - b[K]) N! / (b[1]! b[2]! ... b[K]!) over all possible K-tuples (b[1], b[2], ..., b[K]).

It should be too slow to enumerate all possible K-tuples (b[1], b[2], ..., b[K]) and calculate the corresponding product.

The summation of all the products can be done efficiently with DP. Let f(i, s) be the sum of products of C(N - b[j], a[j] - b[j]) / b[j]! for j = 1, 2, ..., i where b[1], ...., b[i] add up to s. What we need is f(K, N). This DP works in O(N^2 * K). The answer is N! * f(K, N).

This solution works in O(KN^2) time.