SHUFFLEHD - Editorial

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Elementary Combinatorics

PROBLEM:

For a permutation P and a fixed K, define f(P) as the array obtained by doing the following:

  • For i = 1, 2, 3, \ldots N-K+1, sort the subarray P[i, i+K-1] in ascending order, in-place.

Given N, K, and Q, count the number of permutations P for which f(P) = Q.

EXPLANATION:

We solved the easy version by fixing values from left to right; and that reduced nicely since they were 1, 2, \ldots, N in order.

This version can also be solved similarly, by slightly extending the logic.

Let’s look at the value Q_1.
Q_1 must certainly be present among the first K positions of P; otherwise it cannot end up at position 1.
Additionally, Q_1 must be the smallest value among the first K positions of P; since if there’s something smaller than it in that range then upon sorting P[1, K] the first element will be smaller than Q_1, and Q_1 can never move to position 1.

Thus, Q_1 must appear among positions 1, 2, \ldots, K; and everything else in that range must be \gt Q_1.


Next, we look at Q_2.
Again, Q_2 must appear among positions 1, 2, \ldots, K+1.

However, there are now a couple of cases.

First, suppose Q_1 \lt Q_2.
Then, Q_2 can appear at any free position in [1, K+1] (there are K of them), but all remaining free positions are forced to contain values \gt Q_2 only (since if any of them had a smaller value, it’d end up at position 2 after sorting instead.)

On the other hand, if we had Q_1 \gt Q_2 instead, then observe that Q_2 cannot appear at positions 1, 2, \ldots, K at all since we’re constrained by all free positions there containing values larger than Q_1.
This means Q_2 is forced to appear at position K+1 only.

Note that the end result of both cases, however, looks remarkably similar: two positions among 1, 2, \ldots, K+1 are filled; and all free positions in this range have a lower bound on what values they can take (either Q_1 or Q_2, whichever is larger.)
This allows us to generalize the argument.


The generalized property is as follows: for any index x, value Q_x must be placed at a position in 1, 2, \ldots, \min(x+K-1, N).
Additionally, all free positions in this range are forced to contain values greater than \max(Q_1, \ldots, Q_x).

This can be proved inductively: assume it’s true for x-1, and then prove it for x; the proof is similar to what we did for Q_2 and depends on whether Q_x \gt \max(Q_1, \ldots, Q_{x-1}) or not.
In particular, if Q_x is a new maximum then it can be placed in any free position; and if not then it’s forced to be placed in position x+K-1.
Of course, as the base case we know it’s true for x = 1.

This property allows us to count valid permutations.
For each x = 1, 2, \ldots, N,

  • Q_x must be placed within positions 1, 2, \ldots, \min(x+K-1, N).
  • x-1 of these positions are already filled.
  • If Q_x is a new maximum, it can be placed in any free position for \min(x+K-1, N) - (x-1) options.
  • If Q_x is not a new maximum, it can only be placed in position x+K-1, for one option.
    • In particular, note that if x+K-1 \gt N in this case then we don’t actually have anywhere to place Q_x, so the answer is immediately 0.

Multiply all counts to obtain the final answer.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Editorialist's code (PyPy3)
mod = 998244353
for _ in range(int(input())):
    n, k = map(int, input().split())
    q = list(map(int, input().split()))
    
    ans, rem, lo = 1, k-1, 0
    for i in range(1, n+1):
        if i+k-1 <= n: rem += 1
        
        if q[i-1] > lo:
            ans = (ans * rem) % mod
            lo = q[i-1]
        elif i+k-1 > n: ans = 0
        rem -= 1
    print(ans)
2 Likes