PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
Combinatorics
PROBLEM:
For an array A, define f(A) as the minimum number of following operations needed to make all its array equal:
- Add 1 to a prefix of A, or
- Right-rotate the array, i.e. turn it into [A_N, A_1, \ldots, A_{N-1}]
Given N and M, compute the sum of f(A) across all arrays of length N with elements in [1, M].
EXPLANATION:
Let’s understand how to compute f(A) for a fixed array.
Since we’re allowed to rotate the array, we can treat the array as being circular - i.e. A_1 and A_N are adjacent.
In particular, A_{N+1} will refer to A_1, and A_0 will refer to A_N below.
Since we can rotate the array and then choose any prefix, we are essentially able to choose any circular subarray of the circular array A and add 1 to all its elements.
We can use this observation to build a lower bound on the answer.
Consider any index i such that A_i \lt A_{i-1}.
The only way to bring A_i closer to A_{i-1} is to operate on a circular subarray starting at index i.
Thus, we surely need at least A_{i-1} - A_{i} addition operations to begin at this index.
If A_i \ge A_{i-1} we don’t (yet) need to choose any subarrays starting at i.
So, we surely need at least
addition operations to make all the elements equal.
(Note that we are yet to consider rotations.)
It’s not hard to see that this many additions is sufficient: choose any index i such that A_i \lt A_{i-1} and then there will always exist some index j such that A_j \lt A_{j+1} so you can operate on the circular subarray [i, j].
(Quick proof of why such j exists: start at i and keep moving right; you cannot decrease forever because you eventually end up at A_{i-1} which is larger than A_i so you have to increase at some point; choose this as j.)
We thus have a lower bound on the number of addition operations.
Let’s look at the rotations.
With the existing knowledge, this is easy to deal with: we know that only indices with A_i \lt A_{i-1} need to be operated on as the left endpoint of a circular array - which means all such indices must visit index 1 at least once.
Since we right-rotate the array, we go through left endpoints in the order 1, N, N-1, \ldots, 3, 2.
We only need enough right-rotations to get to the last i with A_i \lt A_{i-1} in this order.
In particular, let s be the smallest index in [2, N] such that A_s \lt A_{s-1}, then we need N+1-s rotations.
If no s in [2, N] satisfies this condition, we need 0 rotations.
We now move to summing up this quantity across all arrays.
As seen above, f(A) can be computed by separately summing up additions and rotations.
Let’s first work on additions.
For each index i such that A_i \lt A_{i-1}, we added A_{i-1} - A_i to the answer.
To sum this up across all positions:
- Choose the index i.
N choices here. - Only A_i and A_{i-1} matter, so for each of the remaining N-2 indices, there are M options each for their values.
M^{N-2} choices here. - Finally, we want to sum up (A_{i-1} - A_i) across all choices of A_{i-1} \gt A_i.
- If we fix A_{i-1} = x, then varying A_i from 1 to x-1 gives us a sum of 1 + 2 + \ldots + x-1 = \frac{x\cdot (x-1)}{2}.
- We want to sum this up across all 1 \le x \le M.
Doing the math, this works out to\frac{(M+1)\cdot M \cdot (M-1)}{6}
So we obtain a contribution of
from this part.
Next, we need to sum up the contribution of the rotations.
For that, as we noted, if s denotes the smallest index in [2, N] such that A_s \lt A_{s-1}, then the cost comes out to be N-s+1.
While it’s tempting to try and fix s, it’s not immediately obvious how to do the counting faster than \mathcal{O}(M) after doing that - because we need the prefix upto s-1 to be non-decreasing but A_s \lt A_{s-1}; and counting non-decreasing prefixes is easy but not dealing with the second condition.
Instead, let’s rewrite the criterion slightly.
If s is the smallest index in [2, N] such that A_s \lt A_{s-1}, we’re adding N-s+1; which is equivalent to saying that we add 1 for each index in the range [s, N].
In particular, observe that:
- For each i = 1, 2, \ldots, s-1, the prefix of the array of length i is sorted.
- For each i = s, s+1, \ldots, N, the prefix of the array of length i is not sorted.
So, we’re really just counting the number of not-sorted prefixes of the array!
This observation allows us to do the following.
Fix an index i (2 \le i \le N).
Let’s count the number of arrays in which the prefix of length i is not sorted.
To do this, we instead count the number of arrays in which this prefix is sorted, and subtract from the total.
To count arrays with sorted prefix:
- The prefix has length i, and must be sorted.
- So, only the frequency of elements in this prefix matters; since once the counts of each element are chosen their order is fixed.
- By stars-and-bars, the number of ways is\binom{M+i-1}{i}because we have M non-negative variables (i.e. one variable representing the count of each value 1, 2, \ldots, M) whose sum must equal i.
- As for indices beyond i, the values there don’t matter at all.
There are N-i of them, each with N choices.
Thus, the number of arrays with sorted prefixes of length i equals
The number of arrays with a not sorted prefix of length i is then just obtained by subtracting this from M^N, the total number of arrays.
Thus, the final answer by putting both parts together is simply
This can easily be computed in \mathcal{O}(N\log{MOD}) time, and even \mathcal{O}(N) with some precomputation and extra work.
TIME COMPLEXITY:
\mathcal{O}(N \log{MOD}) per testcase.
CODE:
Editorialist's code (PyPy3)
mod = 998244353
N = 400005
fac = [1] + list(range(N))
for i in range(1, N): fac[i] = fac[i-1] * i % mod
def C(n, r):
if r < 0 or r > n: return 0
return fac[n] * pow(fac[r] * fac[n-r], mod-2, mod) % mod
for _ in range(int(input())):
n, m = map(int, input().split())
ans = n * pow(m, n-2, mod) * m * (m-1) * (m+1) * pow(6, mod-2, mod)
for i in range(2, n+1):
ans += pow(m, n, mod) - C(m+i-1, i) * pow(m, n-i, mod)
print(ans % mod)