MEXMAX7 - Editorial

editorial
, , ,
1

PROBLEM LINK:

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

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Simple

PREREQUISITES:

Combinatorics

PROBLEM:

You’re given an array A.
Count the number of its subsequences S for which |\max(S) - \text{mex}(S)| \le 1.

EXPLANATION:

Suppose we fix M = \max(S) to be the largest element of the subsequence.

Then, we need \text{mex}(S) to be either M-1, M, or M+1.
Note that it can’t be M itself since the mex can’t belong to the sequence by definition (and M must belong to the sequence), so we’re left with only two options.

Let’s try to count valid subsequences with \text{mex}(S) = M-1.
For this to be the case, the subsequence must contain at least one copy of each value 0, 1, 2, \ldots, M-2, then must not contain any copies of M-1, and then must contain at least one copy of M to satisfy \max(S) = M.
Elements larger than M can’t belong to the subsequence.

Counting such subsequences is not hard.
Let’s define f_x to be the number of occurrences of x in the whole array.
Then, observe that there are 2^{f_x} - 1 ways to choose a subset of the occurrences of x, such that at least one of them is chosen.
This is because we have two choices for each occurrence (choose or not), with the only bad choice being when we end up not choosing all of them (so only one possibility.)

Thus, the number of valid subsequences with \text{mex}(S) = M-1 equals

(2^{f_0} - 1)\cdot (2^{f_1} - 1) \cdot (2^{f_2} - 1) \cdot\ldots\cdot (2^{f_{M-2}} - 1) \cdot (2^{f_M} - 1)

By similar reasoning, the number of valid subsequences with \text{mex}(S) = M+1 equals

(2^{f_0} - 1)\cdot (2^{f_1} - 1) \cdot (2^{f_2} - 1) \cdot\ldots\cdot (2^{f_{M-2}} - 1) \cdot (2^{f_{M-1}} - 1)\cdot (2^{f_M} - 1)

that is, essentially the same expression but this time we also need a positive number of occurrences of M-1 as well.

This can be easily computed in \mathcal{O}(N) for a fixed value of M, and trying all M = 0, 1, 2, \ldots, N gives a solution in \mathcal{O}(N^2) time.
It’s possible to improve the complexity to \mathcal{O}(N) by maintaining prefix products, but the constraints are low so this wasn’t needed to get AC.

TIME COMPLEXITY:

\mathcal{O}(N) or \mathcal{O}(N^2) per testcase.

CODE:

Editorialist's code (PyPy3) 
mod = 998244353

for _ in range(int(input())):
    n = int(input())
    a = list(map(int, input().split()))
    
    ways = [1]*(n+1)
    for x in a:
        ways[x] = (2*ways[x]) % mod
    
    ans = 0
    for m in range(0, n+1):
        prod = 1
        for x in range(m-1):
            prod = (prod * (ways[x] - 1)) % mod
        prod = (prod * (ways[m] - 1)) % mod
        
        ans += prod
        if m > 0: ans += prod * (ways[m-1] - 1) % mod
    print(ans % mod)
1 Like