Hints

Hint 1

First, analyze the structure of good arrays.

Hint 2

The sum of elements of a good array has to be not divisible by $P$. Is this enough?

Hint 3

Prove that if every element appears not more than $\frac{N}{2}$ times (and the sum is not divisible by $P$), the array is good.

Hint 4

Consider the most frequent element, let it be $x$. Multiply all elements by $x^{-1} \bmod P$. What can you say about this array now?

Hint 5

Suppose that we have some nonzero elements $A_1, A_2, \dots, A_K$ in our array. At most how many ones can we add to it, so that the array remains good?

Hint 6

If we have some elements $A_1, A_2, \dots, A_K$, not equal to $1$, in our array, we can add at most $(P - A_1) + (P-A_2) + \dots + (P-A_K) + (P-1)$ ones to the array. Now, how to solve the problem?

Hint 7

Write dp

Solution

First of all, let’s figure out when the array is good.

If the sum of all elements is divisible by \(P\), it’s clearly impossible. From now on, we suppose that the sum is not divisible by \(P\).

Let \(X\) be the most frequent element (or one of them), suppose it appears \(M\) times in the array. Multiply all numbers by \(X^{-1}\) so that now \(1\) is the most frequent number.

Let \(B_1, B_2, \dots, B_K\) be all the elements larger than \(1\) in our array.

Then reordering is possible if and only if the number of ones doesn’t exceed \((P - B_1) + (P - B_2) + \dots + (P - B_K) + P-1\).

Proof of necessity

Suppose that the number of ones larger than this, then it's at least $(P - B_1) + (P - B_2) + \dots + (P - B_K) + P+1$, as the total sum can't be divisible by $P$. So, the total sum of all numbers is at least $P(K+1) + 1$.

We have to "jump" over points $P, 2P, \dots, (K+1)P$, and for every such "jump", we need a number larger than $1$. However, we have only $K$ such numbers, and one number can close at most one "jump", so such reordering isn't possible.

Proof of sufficiency

I claim that we can arrange all the numbers with the following process:

Let $x$ be one of the most frequent numbers at the moment and $cur$ be the current sum.

• If $cur + x$ isn't divisible by $P$, append $x$ to the end
• Else take any number which isn't $x$, say $y$, and append $y$, $x$ in this order.

Firstly, in the second case prefix sums won't be divisible by $p$: if $cur + x$ is divisible by $p$, then $cur + y$ and $cur + y + x$ aren't.

Now, when does this algorithm fail? Only when $cur + x$ is divisible by $P$, and there is no other number than $x$, so all the remaining numbers are $x$.

Note that this means that we have at least $2$ numbers $x$ remaining (as otherwise the sum would be divisible by $P$). So, we have $\ge 2$ numbers $x$ and no others.

Note that this means that $x$ was always the most frequent number (we can prove that if some number $y$ was the most frequent at some point, then at no point in future there can exist an element which appears more times than (number of times which $y$ appears $+ 1$)).

So, $x = 1$, and we were always taking $1$ when possible, ending up with only several ones. This means that our sequence looks something like this:

$P-1$ ones, $B_1$, $P - B_1$ ones, $B_2$, $P-B_2$ ones, $\dots$, $B_K$, $P-B_K$ ones, and only now we don't have an option, so we have at least one more one.

But, then the number of ones would exceed the number from our statement, contradiction.

Now we just have to write dp using this knowledge. We will count the number of bad arrays instead.

The number of arrays for which sum is divisible by \(P\) is \(\frac{(P-1)^N - (P-1)}{P}\) for odd \(N\) and \(\frac{(P-1)^N + (P-1)}{P}\) for even \(N\).

Now let’s count the number of arrays which don’t satisfy the second criteria (and have sum not divisible by \(P\) at the same time). We will count only those for which number of ones is too large, and multiply by \(P-1\).

Let’s write \(dp[cnt][sum]\), the number of arrays of \(cnt\) numbers from \(2\) to \(P-1\) such that the sum of \(P - x\) over \(x\) in such array is \(sum\). This can easily be done in \(O(N^2)\). The remaining part is iterating over all pairs \((cnt, sum)\), checking if the number of ones \(N - cnt\) satisfies \(N - cnt \ge sum + P\) and \(N - cnt \neq sum \bmod P\), if so, we add to the answer \(dp[cnt][sum]\binom{N}{cnt}\).

Total asymptotics \(O(N^2)\).