You have to output the sum of the happiness over all possible combinations of choices.

Use the trick. At the end of the process, for each child color one of his cookies in red. You have to output the number of combination of choices, if you assume that the coloring is a choice too. What if you distribute red cookies during the $$$k$$$ days instead of coloring cookies at the end?

Now the problem becomes "In the $$$i$$$-th day, we will choose $$$a_i$$$ children and give either a red cookie or a normal cookie to each child chosen. In how many ways each children ends up having exactly one red cookie?"

Let's find a DP. You can simulate rounds. Which information do you need after each round?

After each round, you only need the number of children who have already received a red cookie.

If you use the trick, the problem becomes "In the $$$i$$$-th day, we will choose $$$a_i$$$ children and give either a red cookie or a normal cookie to each child chosen. In how many ways each children ends up having exactly one red cookie?".

You can calculate the answer with DP. Let $$$dp_{i, j}$$$ be the number of ways to distribute the candies in the first $$$i$$$ rounds and give a red cookie to exactly $$$j$$$ children. When you calculate $$$dp_{i, j}$$$, you can iterate over the number $$$m$$$ of red cookies distributed in the last round. The transition is $$$dp_{i, j} = \sum_{m=0}^{\min(j, a_i)} dp_{i-1, j-m} \cdot \binom{n-j+m}{m} \cdot \binom{n-m}{a_i-m}$$$. In fact, there are

• $$$dp_{i-1, j-m}$$$ ways to distribute candies in the previous rounds
• $$$\binom{n-j+m}{m}$$$ ways to choose the children who receive a red cookie
• $$$\binom{n-m}{a_i-m}$$$ ways to choose the children who receive a normal cookie

The answer is $$$dp_{k, n}$$$.

Complexity: $$$O(n^2k)$$$.

This problem is very similar to the previous one. The differences are that the boxes already contain some balls and $$$k$$$ is large (but the operation are simpler, because you insert exactly $$$1$$$ ball for each operation).

Let's deal with the first difference. In the previous problem, when did we use the fact that the boxes were initially empty?

Hint 5 in the previous problem is true because the boxes were initially empty (so they also were initially equal). In this problem, is it possible to avoid dealing with the $$$a_i$$$ during the operations (e.g. can you do some preprocessing before considering the operations)?

Let's color the balls already in the boxes. You can calculate the number of ways to color $$$m$$$ balls with DP.

Now we can iterate over the number $$$m$$$ of balls colored initially.

For each $$$m$$$, the problem becomes "In how many ways can you put $$$k$$$ balls in the boxes and color $$$n-m$$$ of them, one for each box without already colored balls?"

Instead of coloring the balls at the end, you can choose the operations where you put a red ball in a box. In how many ways can you do that?

Use the trick. First, fix the number of colored balls that were already in the boxes before the operations. Let $$$dp_{i, j}$$$ be the number of ways to color exactly $$$j$$$ balls in the first $$$i$$$ boxes. The transition is $$$dp_{i, j} = dp_{i-1, j} + a_i \cdot dp_{i-1, j-1}$$$.

If you color exactly $$$m$$$ boxes before the operations, the number of ways is $$$dp_{n, m}$$$, multiplied by the number of ways to put $$$k$$$ balls in the boxes and color $$$n-m$$$ of them, one for each box without already colored balls.

Now let's deal with the $$$k$$$ operations. Let's fix which operations add a red ball ($$$\binom{k}{n-m}$$$ ways), the order of red balls ($$$(n-m)!$$$ ways) and the boxes chosen in the other operations ($$$n^{k-n+m}$$$).

Therefore, the answer is $$$\frac{1}{n^m} \cdot \sum_{m=0}^{n} dp_{n, m} \cdot \binom{k}{n-m} \cdot (n-m)! \cdot n^{k-n+m}$$$

Complexity: $$$O(n^2)$$$.

You can replace the DP with a small to large NTT, and the complexity becomes $$$O(n \log^2 n)$$$.

In the official editorial, there is a solution with polynomials and linearity of expectation. Find the similarities between this solution and the one in the editorial.

$$$S$$$ can be obtained by only considering sequences where at least a person hands $$$0$$$ balls to the neighbor on his right.

Use the trick. For each person $$$i$$$, his red ball was initially owned either by $$$i$$$ or by the neighbor on his left.

Use DP. In how many ways can you distribute and color the balls for each prefix? You need additional information:

• the array is circular, so you need to know if person $$$1$$$ has already a colored ball (that wasn't given by person $$$n$$$)
• while processing person $$$i$$$, you need to know if he has already received a colored ball
• you need to know if at least a person in the prefix has given $$$0$$$ balls to the neighbor on his right

Let's do an example of transition. If you want to color one of the balls of person $$$i$$$ and give some balls to his neighbor, including a colored ball, the number of ways can be calculated by fixing the number $$$m$$$ of balls given. For a fixed $$$m$$$, the answer is $$$m(a_i-m)$$$. The total number of ways is $$$\sum_{m=1}^{a_i} m(a_i-m) = a_i \cdot \sum_{m=1}^{a_i} m - \sum_{m=1}^{a_i} m^2 = \frac{a_i^2(a_i+1)}{2} - \frac{a_i(a_i+1)(2a_i+1)}{6}$$$. The other transitions are similar.

Use the trick. Let $$$dp_{i, j, k, l}$$$ the number of ways the people numbered from $$$1$$$ to $$$i$$$ can give and color balls, considering that

• the person $$$1$$$ has colored $$$j$$$ balls ($$$0 \leq j \leq 1$$$)
• the neighbor to the right of person $$$i$$$ has received $$$k$$$ colored balls ($$$0 \leq k \leq 1$$$)
• if $$$l = 0$$$, no person has given $$$0$$$ balls to his neighbor; if $$$l = 1$$$, at least one person has given $$$0$$$ balls to his neighbor.

The transitions can be found by iterating on the answers of the following questions:

• must the person $$$i$$$ color one of his balls?
• how many balls can $$$i$$$ give to his neighbor?
• must the person $$$i$$$ give a colored ball to his neighbor?

An example of transition (with formulas) is in Hint 4.

The answer is the sum of $$$dp_{n, j, 1-j, 1}$$$ (because person $$$1$$$ must have exactly $$$1$$$ colored ball, and there must be a person who has given $$$0$$$ balls to his neighbor).

Complexity: $$$O(n)$$$.