Tokitsukaze and Permutations (Codeforces 1677D)
Tags: Math, dp

Tokitsukaze has a permutation $$$p$$$. She performed the following operation to $$$p$$$ exactly $$$k$$$ times: in one operation, for each $$$i$$$ from $$$1$$$ to $$$n - 1$$$ in order, if $$$p_i$$$ > $$$p_{i+1}$$$, swap $$$p_i$$$, $$$p_{i+1}$$$. After exactly $$$k$$$ times of operations, Tokitsukaze got a new sequence $$$a$$$, obviously the sequence $$$a$$$ is also a permutation.After that, Tokitsukaze wrote down the value sequence $$$v$$$ of $$$a$$$ on paper. Denote the value sequence $$$v$$$ of the permutation $$$a$$$ of length $$$n$$$ as $$$v_i=\sum_{j=1}^{i-1}[a_i < a_j]$$$, where the value of $$$[a_i < a_j]$$$ define as if $$$a_i < a_j$$$, the value is $$$1$$$, otherwise is $$$0$$$ (in other words, $$$v_i$$$ is equal to the number of elements greater than $$$a_i$$$ that are to the left of position $$$i$$$). Then Tokitsukaze went out to work.There are three naughty cats in Tokitsukaze's house. When she came home, she found the paper with the value sequence $$$v$$$ to be bitten out by the cats, leaving several holes, so that the value of some positions could not be seen clearly. She forgot what the original permutation $$$p$$$ was. She wants to know how many different permutations $$$p$$$ there are, so that the value sequence $$$v$$$ of the new permutation $$$a$$$ after exactly $$$k$$$ operations is the same as the $$$v$$$ written on the paper (not taking into account the unclear positions).Since the answer may be too large, print it modulo $$$998\,244\,353$$$.

Input: The first line contains a single integer $$$t$$$ ($$$1 \leq t \leq 1000$$$) — the number of test cases. Each test case consists of two lines.The first line contains two integers $$$n$$$ and $$$k$$$ ($$$1 \leq n \leq 10^6$$$; $$$0 \leq k \leq n-1$$$) — the length of the permutation and the exactly number of operations.The second line contains $$$n$$$ integers $$$v_1, v_2, \dots, v_n$$$ ($$$-1 \leq v_i \leq i-1$$$) — the value sequence $$$v$$$. $$$v_i = -1$$$ means the $$$i$$$-th position of $$$v$$$ can't be seen clearly.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$10^6$$$.

Output: For each test case, print a single integer — the number of different permutations modulo $$$998\,244\,353$$$.

Note: In the first test case, only permutation $$$p=[5,4,3,2,1]$$$ satisfies the constraint condition.In the second test case, there are $$$6$$$ permutations satisfying the constraint condition, which are:  $$$[3,4,5,2,1]$$$ $$$\rightarrow$$$ $$$[3,4,2,1,5]$$$ $$$\rightarrow$$$ $$$[3,2,1,4,5]$$$  $$$[3,5,4,2,1]$$$ $$$\rightarrow$$$ $$$[3,4,2,1,5]$$$ $$$\rightarrow$$$ $$$[3,2,1,4,5]$$$  $$$[4,3,5,2,1]$$$ $$$\rightarrow$$$ $$$[3,4,2,1,5]$$$ $$$\rightarrow$$$ $$$[3,2,1,4,5]$$$  $$$[4,5,3,2,1]$$$ $$$\rightarrow$$$ $$$[4,3,2,1,5]$$$ $$$\rightarrow$$$ $$$[3,2,1,4,5]$$$  $$$[5,3,4,2,1]$$$ $$$\rightarrow$$$ $$$[3,4,2,1,5]$$$ $$$\rightarrow$$$ $$$[3,2,1,4,5]$$$  $$$[5,4,3,2,1]$$$ $$$\rightarrow$$$ $$$[4,3,2,1,5]$$$ $$$\rightarrow$$$ $$$[3,2,1,4,5]$$$ So after exactly $$$2$$$ times of swap they will all become $$$a=[3,2,1,4,5]$$$, whose value sequence is $$$v=[0,1,2,0,0]$$$.

Sample Inputs:
3
5 0
0 1 2 3 4
5 2
-1 1 2 0 0
5 2
0 1 1 0 0

Sample Outputs:
1
6
6

Solution (Python):
mod = 998244353
t = int(input())
for _ in range(t):
    (n, k) = map(int, input().split(' '))
    f = 0
    v = list(map(int, input().split()))
    value = [0]
    value.extend(v)
    ans = 1
    for i in range(1, k + 1):
        ans = ans * i % mod
    for i in range(n - k + 1, n + 1):
        if value[i] != -1 and value[i]:
            print(0)
            f = 1
            break
    if f:
        continue
    for i in range(k + 1, n + 1):
        limit = i - 1
        if value[i - k] == -1:
            ans = ans * i % mod
        elif value[i - k]:
            limit = value[i - k] + k
            if limit > i - 1:
                f = 1
        else:
            limit = min(limit, k)
            ans = ans * (limit + 1) % mod
    if f:
        ans = 0
    print(ans)

Solution (Python):
import sys
I = lambda : [*map(int, sys.stdin.readline().split())]
M = 998244353
facs = [1]
for i in range(1000000):
    facs.append(facs[-1] * (i + 1) % M)
(t,) = I()
for _ in range(t):
    (n, k) = I()
    p = I()
    bad = False
    for i in range(n):
        if p[i] > i:
            bad = True
        if i >= n - k and p[i] > 0:
            bad = True
    if bad:
        print(0)
        continue
    inv = pow(k + 1, M - 2, M)
    stuff = [facs[k] * pow(k + 1, n - k, M) % M]
    for i in range(n):
        stuff.append(stuff[-1] * inv % M)
    big = 0
    for guy in p:
        if guy > 0:
            big += 1
    tot = stuff[big]
    for i in range(n - k):
        if p[i] == -1:
            poss = i
            tot = (tot + poss * tot * inv) % M
    print(tot)

Solution (Python):
MOD = 998244353
t = int(input())
for _ in range(t):
    (n, k) = map(int, input().split())
    v = list(map(int, input().split()))
    out = 1
    for i in range(n - k):
        if v[i] == 0:
            out *= k + 1
            out %= MOD
        if v[i] == -1:
            out *= i + k + 1
            out %= MOD
    for i in range(k):
        if v[~i] > 0:
            out = 0
        out *= i + 1
        out %= MOD
    print(out)