Official

F - Creative Splitting Editorial by antontrygubO_o


Let’s first design a greedy algorithm to determine whether the sequence is amazing. We will maintain \(N\) sequences \(S_1, S_2, \ldots, S_N\), and will construct them from the end. For every element from \(A_{NK}\) to \(A_1\), we will assign it to the longest among the sequences to which we can assign it. Then, if for some element there weren’t any subsequences to assign it to, there is no solution, otherwise, we have already constructed the required split.

Let’s show that if it’s possible to split the sequence as the statement wants us to, then the algorithm above will also produce the correct split. Again, let’s process the elements from the end, and assign them according to the “correct” split. Clearly, it’s possible to assign element \(X\) to a subsequence \(S_i\) if and only if currently \(X \le K - |S_i|\). Consider the last moment when the element \(A_i\) got assigned to the subsequence \(S_x\), but there exists a subsequence \(S_y\) with \(|S_x|<|S_y|\le K-A_i\), choose such \(S_y\) with the largest \(|S_y|\).

Then consider the next assigned element to \(S_y\), \(A_j\), with \(j<i\). Clearly, we won’t assign any of the elements between \(A_{j+1}\) and \(A_{i-1}\) to \(S_x\), as we have chosen the last such “bad” moment. Then, it turns out that we could have just assigned \(A_i\) to \(S_y\) and \(A_j\) to \(S_x\) instead. Indeed, this would work, as we have \(A_i \le K - |S_y|\) and \(A_j \le K - |S_y| < K - |S_x|\).

It’s easy to see that this process is finite. For example, we can prove it as follows: consider a binary string \(T_{NK}T_{NK-1}\ldots T_1\), where \(T_i\) is \(0\) if \(i\)-th element wasn’t assigned to the longest among available subsequences, and \(1\) otherwise, then during this process the string always increases lexicographically, and the greedy algorithm is proved.

Now, consider the following process. Consider a strip of length \(N+1\), with cells numbered from \(0\) to \(N\), with \(K\) balls lying in the first cell. In one operation, we can choose any ball which isn’t in the cell \(N\), and move it to the next cell to the right. The process will end after exactly \(NK\) operations.

Surprisingly, there is a bijection between sequences of operations in this process and the amazing arrays! To see it, let’s look at the structure of our greedy a bit closer.

For each sequence, keep track of how much more elements we need to add to it — \(L_i = K - |S_i|\). Initially all \(L_i\) are equal to \(K\), and we want all of them to become \(0\) in the end. Let’s sort the sequences by \(L_i\), wlog \(L_1 \le L_2 \le \ldots \le L_N\). Then, if the current element is \(X\), according to the greedy we will take the smallest \(L_i\) which is at least \(X\), and decrease it by \(1\).

Now let’s look at the process with the balls. Suppose that currently there are \(cnt_i\) balls at cell \(i\) for each \(i\) from \(0\) to \(N-1\). Then consider \(pref_i = cnt_0 + cnt_1 + \ldots + cnt_{i-1}\) for all \(i\) from \(1\) to \(N\). It’s easy to see that \(pref_1 \le pref_2 \le \ldots \le pref_N\) and that the only thing that operation of moving some ball from cell \(X\) to cell \(X+1\) does is decreasing \(pref_{X+1}\) by \(1\). Clearly, initially all \(pref_i\) are \(K\), and we want all of them to become \(0\) in the end.

If you have already started to suspect something at this point, here is one more “coincidence”. How many ways are there to move from the given sequence \(L_1 \le L_2 \ldots \le L_Y < L_{Y+1} \le \ldots \le L_N\) to \(L_1 \le L_2 \ldots \le L_Y < L_{Y+1}-1 \le \ldots \le L_N\)? We will do this only if the new element is in range \([L_Y+1, L_{Y+1}]\), so there are precisely \(L_{Y+1} - L_Y\) ways to do this.

And how many ways are there to move from \(pref_1 \le pref_2 \ldots \le pref_Y < pref_{Y+1} \le \ldots \le pref_N\) to \(pref_1 \le pref_2 \ldots \le pref_Y < pref_{Y+1} - 1 \le \ldots \le pref_N\)? It’s equal to the number of balls in the cell \(Y\), which is just equal to \(pref_{Y+1} - pref_Y\)!

So, we clearly have established a bijection between the sequences \((L_1, \ldots, L_N)\) and \((pref_1, \ldots, pref_N)\) in our processes. Time to reformulate our problem from the language of amazing arrays to the language of strip and balls.

\(f(pos, val)\) will then be equal to the number of sequences of operations, such that on the \(NK-pos+1\)-st operation we will move a ball from \(X\) to \(X+1\), where \(pref_X < val \le pref_{X+1}\). Let’s solve this problem for all pairs \((pos, val)\).

Let’s suppose that on the \(i\)-th operation we are moving the ball from cell \(pos\) to \(pos+1\). Suppose that \(j\)-th of the other balls is currently at position \(pos_i\). Clearly, \(pos + pos_1 + pos_2 + \ldots + pos_{K-1} = i-1\) must hold. Note that moving all balls is independent of each other. That means that the number of ways to arrive at the current situation is \(\frac{(i-1)!}{pos! pos_1! pos_2! \ldots pos_{K-1}!}\), and the number of ways to get from it (after performing this \(i\)-th operation) is \(\frac{(NK-i)!}{(N-pos-1)!(N-pos_1)!\ldots (N-pos_{K-1})!}\). So, we have to find the sum of \(\frac{(i-1)!}{pos! pos_1! pos_2! \ldots pos_{K-1}!} \cdot \frac{(NK-i)!}{(N-pos-1)!(N-pos_1)!\ldots (N-pos_{K-1})!}\) over all \((i, pos, pos_1, pos_2, \ldots, pos_{N-1})\) such that \(pref_{pos_1} < val < pref_{pos_1 + 1}\). To do this, we also need to take care of the number of balls to the left and to the right of cell \(pos\).

We can do this with some precalculation. Indeed, for a ball its contribution to the final product is \(g_{pos} = \frac{1}{pos!(N-pos)!}\), so for every triple \((num, bound, sum)\) let’s count the sum of \(\prod_{0 \le i < bound} g_{pos}^{cnt_i}\) over all possible assignments of positions of \(num\) balls among the first \(bound\) positions, for which \(\sum_{0 \le i < bound} cnt_i = num\) and \(\sum_{0 \le i < bound} cnt_i\cdot i = sum\). This can be done in \(O(N^2K^2)\) easily.

After that, we go through all possible of \(5\)-tuples of \((pos, num_{left}, num_{right}, sum_{left}, sum_{right}\), and calculate the corresponding number of ways to get to the situation where we move a ball from cell \(pos\) to \(pos+1\) and then to get all balls to the end. With the precalculations above, the complexity becomes \(O(N^3K^4)\) with an incredibly small constant, which passes easily under given limitations.

posted:
last update: