Recently I came across a question in which we had to find the no. of ways to distribute N balls into 3 boxes such that the number of box getting maximum no. of balls is exactly 1.
For example
if N=2, ans is 2 : {2,0,0},{0,2,0},{0,0,2}
if N=3, ans is 9 : {3,0,0},{0,3,0},{0,0,3},{3,2,1},{3,1,2},{1,3,2},{2,3,1},{1,2,3},{2,1,3}
Now, since the number of boxes here is only 3, I was able to solve the question by observation and basic maths(Arith. Prog.) My query is how to solve the above question if the number of boxes is K?
This can be calculated in $$$O((N + K) \log N)$$$. Let $$$\binom{a}{b} = 0$$$ when $$$a < b$$$. Using inclusion-exclusion, we get that the answer is
$$$K \sum_{s = 1}^{N} \sum_{c = 0}^{K-1} (-1)^{c} \binom{K-1}{c} \binom{K-1 + N-s - sc}{K-1}$$$
since $$$K$$$ is the number of ways to select the maximum element, we loop over $$$s$$$ which will be its value, and
$$$\sum_{c = 0}^{K-1} (-1)^{c} \binom{K-1}{c} \binom{K-1 + N-s - sc}{K-1}$$$
calculates the number of ways to select values for the other numbers such that no other number is $$$\geqslant s$$$. Changing the order of summation we get
$$$\sum_{c = 0}^{K-1} (-1)^{c} K \binom{K-1}{c} \sum_{s = 1}^{\left\lfloor\frac{N}{c+1}\right\rfloor} \binom{K-1 + N - s(c+1)}{K-1}$$$
Which can be directly calculated in $$$O((N + K) \log N)$$$ by precalculating factorials to compute the binomial terms in $$$O(1)$$$
Hey, isn’t it problem from ongoing contest? Why you didn’t asked after end of the contest?