PermutationForces II (Codeforces 1698E)
Tags: Greedy, combinatorics

You are given a permutation $$$a$$$ of length $$$n$$$. Recall that permutation is an array consisting of $$$n$$$ distinct integers from $$$1$$$ to $$$n$$$ in arbitrary order.You have a strength of $$$s$$$ and perform $$$n$$$ moves on the permutation $$$a$$$. The $$$i$$$-th move consists of the following:   Pick two integers $$$x$$$ and $$$y$$$ such that $$$i \leq x \leq y \leq \min(i+s,n)$$$, and swap the positions of the integers $$$x$$$ and $$$y$$$ in the permutation $$$a$$$. Note that you can select $$$x=y$$$ in the operation, in which case no swap will occur. You want to turn $$$a$$$ into another permutation $$$b$$$ after $$$n$$$ moves. However, some elements of $$$b$$$ are missing and are replaced with $$$-1$$$ instead. Count the number of ways to replace each $$$-1$$$ in $$$b$$$ with some integer from $$$1$$$ to $$$n$$$ so that $$$b$$$ is a permutation and it is possible to turn $$$a$$$ into $$$b$$$ with a strength of $$$s$$$. Since the answer can be large, output it modulo $$$998\,244\,353$$$.

Input: The input consists of multiple test cases. The first line contains an integer $$$t$$$ ($$$1 \leq t \leq 1000$$$) — the number of test cases. The description of the test cases follows.The first line of each test case contains two integers $$$n$$$ and $$$s$$$ ($$$1 \leq n \leq 2 \cdot 10^5$$$; $$$1 \leq s \leq n$$$) — the size of the permutation and your strength, respectively.The second line of each test case contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$1 \le a_i \le n$$$) — the elements of $$$a$$$. All elements of $$$a$$$ are distinct.The third line of each test case contains $$$n$$$ integers $$$b_1, b_2, \ldots, b_n$$$ ($$$1 \le b_i \le n$$$ or $$$b_i = -1$$$) — the elements of $$$b$$$. All elements of $$$b$$$ that are not equal to $$$-1$$$ are distinct.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$2 \cdot 10^5$$$.

Output: For each test case, output a single integer — the number of ways to fill up the permutation $$$b$$$ so that it is possible to turn $$$a$$$ into $$$b$$$ using a strength of $$$s$$$, modulo $$$998\,244\,353$$$.

Note: In the first test case, $$$a=[2,1,3]$$$. There are two possible ways to fill out the $$$-1$$$s in $$$b$$$ to make it a permutation: $$$[3,1,2]$$$ or $$$[3,2,1]$$$. We can make $$$a$$$ into $$$[3,1,2]$$$ with a strength of $$$1$$$ as follows: $$$$$$[2,1,3] \xrightarrow[x=1,\,y=1]{} [2,1,3] \xrightarrow[x=2,\,y=3]{} [3,1,2] \xrightarrow[x=3,\,y=3]{} [3,1,2].$$$$$$ It can be proven that it is impossible to make $$$[2,1,3]$$$ into $$$[3,2,1]$$$ with a strength of $$$1$$$. Thus only one permutation $$$b$$$ satisfies the constraints, so the answer is $$$1$$$.In the second test case, $$$a$$$ and $$$b$$$ the same as the previous test case, but we now have a strength of $$$2$$$. We can make $$$a$$$ into $$$[3,2,1]$$$ with a strength of $$$2$$$ as follows: $$$$$$[2,1,3] \xrightarrow[x=1,\,y=3]{} [2,3,1] \xrightarrow[x=2,\,y=3]{} [3,2,1] \xrightarrow[x=3,\,y=3]{} [3,2,1].$$$$$$ We can still make $$$a$$$ into $$$[3,1,2]$$$ using a strength of $$$1$$$ as shown in the previous test case, so the answer is $$$2$$$. In the third test case, there is only one permutation $$$b$$$. It can be shown that it is impossible to turn $$$a$$$ into $$$b$$$, so the answer is $$$0$$$.

Sample Inputs:
6
3 1
2 1 3
3 -1 -1
3 2
2 1 3
3 -1 -1
4 1
1 4 3 2
4 3 1 2
6 4
4 2 6 3 1 5
6 1 5 -1 3 -1
7 4
1 3 6 2 7 4 5
2 5 -1 -1 -1 4 -1
14 14
1 2 3 4 5 6 7 8 9 10 11 12 13 14
-1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1

Sample Outputs:
1
2
0
2
12
331032489

Solution (Python):
t = int(input())
for _ in range(t):
    (n, s) = map(int, input().split())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))
    visited = [0] * (n + 1)
    m = 1
    empty = []
    for (a_i, b_i) in zip(a, b):
        if b_i != -1 and a_i - b_i > s:
            m = 0
            break
        elif b_i == -1:
            empty.append(a_i)
        else:
            visited[b_i] += 1
    else:
        missing = [i for i in range(1, n + 1) if visited[i] == 0]
        empty.sort()
        l_max = len(missing)
        e_hi = 0
        for (e_lo, m_i) in enumerate(missing):
            while e_hi != l_max and m_i + s >= empty[e_hi]:
                e_hi += 1
            m = m * max(0, e_hi - e_lo) % 998244353
    print(m)

Solution (Python):
from functools import reduce
from sys import stdin, stdout

def read():
    return stdin.readline().strip()

def write(x):
    stdout.write(str(x) + '\n')

def get_index(a):
    a_index = [-1] * len(a)
    for i in range(len(a)):
        if a[i] != -1:
            a_index[a[i]] = i
    return a_index

def solve():
    (n, s) = map(int, read().split())
    a_ = list(map(int, read().split()))
    b_ = list(map(int, read().split()))
    a = [i - 1 for i in a_]
    b = [i - 1 if i != -1 else -1 for i in b_]
    a_index = get_index(a)
    b_index = get_index(b)
    a = [a_index[i] for i in range(n)]
    b = [b_index[i] if b_index[i] != -1 else -1 for i in range(n)]
    a_index = get_index(a)
    b_index = get_index(b)
    p_in_b = [a[i] for i in range(s) if b_index[a[i]] != -1]
    p_not_in_b = [a[i] for i in range(s) if b_index[a[i]] == -1]
    p_in_b_index = [-1] * n
    for i in range(len(p_in_b)):
        p_in_b_index[p_in_b[i]] = i
    c = [0] * n
    for i in range(n):
        j = i + s
        if j < n:
            if b_index[a[j]] != -1:
                p_in_b.append(a[j])
                p_in_b_index[a[j]] = len(p_in_b) - 1
            else:
                p_not_in_b.append(a[j])
        if b[i] == -1:
            c[i] = len(p_not_in_b)
            if p_not_in_b:
                p_not_in_b.pop()
        elif p_in_b_index[b[i]] != -1:
            c[i] = 1
            p_in_b_index[b[i]] = -1
        else:
            c[i] = 0
            break
    write(reduce(lambda x, y: x * y % 998244353, c))
T = 1
T = int(read())
for _ in range(T):
    solve()

Solution (Python):
def convert(nums_a, nums_b, s):
    zips = list([(a, b) for (a, b) in zip(nums_a, nums_b)])
    perm = [x[1] for x in sorted(zips)]
    free_nums = set(nums_a) - set(perm)
    available_nums = set(nums_a)
    count_ar = [0] * len(perm)
    count = 1
    for (i, v) in enumerate(perm):
        rm_num = i + 1 - s - 1
        if rm_num in free_nums:
            free_nums.remove(rm_num)
        if rm_num in available_nums:
            available_nums.remove(rm_num)
        count_ar[i] = len(free_nums)
    rm_count = 0
    for (i, v) in reversed(list(enumerate(perm))):
        if perm[i] == -1:
            count *= count_ar[i] - rm_count
            count = count % 998244353
            rm_count += 1
        elif i - perm[i] >= s:
            count = 0
    print(count)
t = int(input())
for _ in range(t):
    (_, s) = list(map(int, input().split()))
    nums_a = list(map(int, input().split()))
    nums_b = list(map(int, input().split()))
    convert(nums_a, nums_b, s)