Counting Prefixes (Codeforces 1919E)
Tags: Math, combinatorics

There is a hidden array $$$a$$$ of size $$$n$$$ consisting of only $$$1$$$ and $$$-1$$$. Let $$$p$$$ be the prefix sums of array $$$a$$$. More formally, $$$p$$$ is an array of length $$$n$$$ defined as $$$p_i = a_1 + a_2 + \ldots + a_i$$$. Afterwards, array $$$p$$$ is sorted in non-decreasing order. For example, if $$$a = [1, -1, -1, 1, 1]$$$, then $$$p = [1, 0, -1, 0, 1]$$$ before sorting and $$$p = [-1, 0, 0, 1, 1]$$$ after sorting.You are given the prefix sum array $$$p$$$ after sorting, but you do not know what array $$$a$$$ is. Your task is to count the number of initial arrays $$$a$$$ such that the above process results in the given sorted prefix sum array $$$p$$$. As this number can be large, you are only required to find it modulo $$$998\,244\,353$$$.

Input: Each test contains multiple test cases. The first line contains a single 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 a single integer $$$n$$$ ($$$1 \le n \le 5000$$$) — the size of the hidden array $$$a$$$.The second line of each test case contains $$$n$$$ integers $$$p_1, p_2, \ldots, p_n$$$ ($$$|p_i| \le n$$$) — the $$$n$$$ prefix sums of $$$a$$$ sorted in non-decreasing order.It is guaranteed that $$$p_1 \le p_2 \le \ldots \le p_n$$$.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$5000$$$.

Output: For each test case, output the answer modulo $$$998\,244\,353$$$.

Note: In the first two test cases, the only possible arrays $$$a$$$ for $$$n = 1$$$ are $$$a = [1]$$$ and $$$a = [-1]$$$. Their respective sorted prefix sum arrays $$$p$$$ are $$$p = [1]$$$ and $$$p = [-1]$$$. Hence, there is no array $$$a$$$ that can result in the sorted prefix sum array $$$p = [0]$$$ and there is exactly $$$1$$$ array $$$a$$$ that can result in the sorted prefix sum array $$$p = [1]$$$.In the third test case, it can be proven that there is no array $$$a$$$ that could result in the sorted prefix sum array $$$p = [-1, 1, 2]$$$.In the fourth test case, the $$$3$$$ possible arrays $$$a$$$ that could result in the sorted prefix sum array $$$p = [-1, 0, 0, 1, 1]$$$ are:  $$$a = [1, -1, 1, -1, -1]$$$. The prefix sum array before sorting is $$$p = [1, 0, 1, 0, -1]$$$, which after sorting gives $$$p = [-1, 0, 0, 1, 1]$$$.  $$$a = [1, -1, -1, 1, 1]$$$. The prefix sum array before sorting is $$$p = [1, 0, -1, 0, 1]$$$, which after sorting gives $$$p = [-1, 0, 0, 1, 1]$$$.  $$$a = [-1, 1, 1, -1, 1]$$$. The prefix sum array before sorting is $$$p = [-1, 0, 1, 0, 1]$$$, which after sorting gives $$$p = [-1, 0, 0, 1, 1]$$$. For the fifth test case, the only possible array $$$a$$$ that could result in the sorted prefix sum array $$$p = [-4, -3, -3, -2, -1]$$$ is $$$a = [-1, -1, -1, -1, 1]$$$.

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

Sample Outputs:
0
1
0
3
1

Solution (Python):
mod = 998244353
fact = [1] * 5005
rfact = [1] * 5005
for i in range(1, 5005):
    fact[i] = fact[i - 1] * i % mod
    rfact[i] = pow(fact[i], mod - 2, mod)

def C(n, r):
    if n < r or r < 0:
        return 0
    return fact[n] * rfact[r] * rfact[n - r] % mod
t = int(input())
for _ in range(t):
    n = int(input())
    arr = list(map(int, input().split())) + [0]
    arr.sort()
    (m, M) = (arr[0], arr[-1])
    flag = 0
    for i in range(n - 1):
        if arr[i + 1] - arr[i] > 1:
            flag = 1
            break
    if flag == 1 or m * M > 0:
        print(0)
    else:
        cnt = [0] * 10005
        for i in arr:
            cnt[i] += 1
        dp = [[[0, 0] for _ in range(cnt[i + m] + 1)] for i in range(M - m + 1)]
        dp[0][cnt[m] - 1][0] = 1
        for i in range(1, M - m + 1):
            p = i + m
            for j in range(cnt[p] + 1):
                for j2 in range(cnt[p - 1] + 1):
                    if p <= 0:
                        hole = cnt[p] - j
                        if hole == j2 + 2:
                            dp[i][j][0] = (dp[i][j][0] + dp[i - 1][j2][0] * C(cnt[p] - 1, hole - 1)) % mod
                        if hole == j2 + 1:
                            dp[i][j][1] = (dp[i][j][1] + dp[i - 1][j2][0] * C(cnt[p] - 1, hole - 1) + dp[i - 1][j2][1] * C(cnt[p] - 1, hole - 1)) % mod
                    else:
                        hole = cnt[p] - j
                        if hole == j2 + 1:
                            dp[i][j][0] = (dp[i][j][0] + dp[i - 1][j2][0] * C(cnt[p] - 1, hole - 1)) % mod
                        if hole == j2:
                            dp[i][j][1] = (dp[i][j][1] + dp[i - 1][j2][0] * C(cnt[p] - 1, hole - 1) + dp[i - 1][j2][1] * C(cnt[p] - 1, hole - 1)) % mod
        ans = sum(dp[-1][0]) % mod
        print(ans)

Solution (Python):
import math
import heapq
import itertools
import bisect
import random
import time
from collections import deque
import sys
from cmath import exp, pi
from functools import cmp_to_key
custom_input = sys.stdin.readline
custom_mod = 998244353
custom_fact = [1] * 5005
custom_rfact = [1] * 5005
for i in range(1, 5005):
    custom_fact[i] = custom_fact[i - 1] * i % custom_mod
    custom_rfact[i] = pow(custom_fact[i], custom_mod - 2, custom_mod)

def custom_combination(n, r):
    if n < r or r < 0:
        return 0
    return custom_fact[n] * custom_rfact[r] * custom_rfact[n - r] % custom_mod
custom_count = [0] * 10005
t_cases = int(custom_input())
for _ in range(t_cases):
    num_elements = int(custom_input())
    custom_array = list(map(int, custom_input().split())) + [0]
    custom_array.sort()
    (min_val, max_val) = (custom_array[0], custom_array[-1])
    flag = 0
    for i in range(num_elements - 1):
        if custom_array[i + 1] - custom_array[i] > 1:
            flag = 1
            break
    if flag == 1 or min_val * max_val > 0:
        print(0)
    else:
        for i in range(min_val, max_val + 1):
            custom_count[i] = 0
        for i in custom_array:
            custom_count[i] += 1
        custom_dp = [[[0, 0] for _ in range(custom_count[i + min_val] + 1)] for i in range(max_val - min_val + 1)]
        custom_dp[0][custom_count[min_val] - 1][0] = 1
        for i in range(1, max_val - min_val + 1):
            p = i + min_val
            for j in range(custom_count[p] + 1):
                for j2 in range(custom_count[p - 1] + 1):
                    if p <= 0:
                        hole = custom_count[p] - j
                        if hole == j2 + 2:
                            custom_dp[i][j][0] = (custom_dp[i][j][0] + custom_dp[i - 1][j2][0] * custom_combination(custom_count[p] - 1, hole - 1)) % custom_mod
                        if hole == j2 + 1:
                            custom_dp[i][j][1] = (custom_dp[i][j][1] + custom_dp[i - 1][j2][0] * custom_combination(custom_count[p] - 1, hole - 1) + custom_dp[i - 1][j2][1] * custom_combination(custom_count[p] - 1, hole - 1)) % custom_mod
                    else:
                        hole = custom_count[p] - j
                        if hole == j2 + 1:
                            custom_dp[i][j][0] = (custom_dp[i][j][0] + custom_dp[i - 1][j2][0] * custom_combination(custom_count[p] - 1, hole - 1)) % custom_mod
                        if hole == j2:
                            custom_dp[i][j][1] = (custom_dp[i][j][1] + custom_dp[i - 1][j2][0] * custom_combination(custom_count[p] - 1, hole - 1) + custom_dp[i - 1][j2][1] * custom_combination(custom_count[p] - 1, hole - 1)) % custom_mod
        result = 0
        for i in custom_dp[-1][0]:
            result = (result + i) % custom_mod
        print(result)

Solution (Python):
import math
import heapq
import itertools
import bisect
import random
import time
from collections import deque
import sys
from cmath import exp, pi
from functools import cmp_to_key
input = sys.stdin.readline
mod = 998244353
fact = [1] * 5005
rfact = [1] * 5005
for i in range(1, 5005):
    fact[i] = fact[i - 1] * i % mod
    rfact[i] = pow(fact[i], mod - 2, mod)

def C(n, r):
    if n < r or r < 0:
        return 0
    return fact[n] * rfact[r] * rfact[n - r] % mod
cnt = [0] * 10005
t = int(input())
for _ in range(t):
    n = int(input())
    arr = list(map(int, input().split())) + [0]
    arr.sort()
    (m, M) = (arr[0], arr[-1])
    flag = 0
    for i in range(n - 1):
        if arr[i + 1] - arr[i] > 1:
            flag = 1
            break
    if flag == 1 or m * M > 0:
        print(0)
    else:
        for i in range(m, M + 1):
            cnt[i] = 0
        for i in arr:
            cnt[i] += 1
        dp = [[[0, 0] for _ in range(cnt[i + m] + 1)] for i in range(M - m + 1)]
        dp[0][cnt[m] - 1][0] = 1
        for i in range(1, M - m + 1):
            p = i + m
            for j in range(cnt[p] + 1):
                for j2 in range(cnt[p - 1] + 1):
                    if p <= 0:
                        hole = cnt[p] - j
                        if hole == j2 + 2:
                            dp[i][j][0] = (dp[i][j][0] + dp[i - 1][j2][0] * C(cnt[p] - 1, hole - 1)) % mod
                        if hole == j2 + 1:
                            dp[i][j][1] = (dp[i][j][1] + dp[i - 1][j2][0] * C(cnt[p] - 1, hole - 1) + dp[i - 1][j2][1] * C(cnt[p] - 1, hole - 1)) % mod
                    else:
                        hole = cnt[p] - j
                        if hole == j2 + 1:
                            dp[i][j][0] = (dp[i][j][0] + dp[i - 1][j2][0] * C(cnt[p] - 1, hole - 1)) % mod
                        if hole == j2:
                            dp[i][j][1] = (dp[i][j][1] + dp[i - 1][j2][0] * C(cnt[p] - 1, hole - 1) + dp[i - 1][j2][1] * C(cnt[p] - 1, hole - 1)) % mod
        ans = 0
        for i in dp[-1][0]:
            ans = (ans + i) % mod
        print(ans)