Subset MEX Counting Practice Coding Problem [Link]

Subset MEX Counting Practice Coding Problem

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

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:
Easy

PREREQUISITES:
Combinatorics, prefix sums

For an array 𝐴 and an integer 𝑀, define 𝑓 (𝐴, 𝑀) to be the number of distinct MEX-es across all size 𝑀
subsets of 𝐴.

You’re given an integer 𝑁.

For each 𝑀 = 1, 2, 3, …, 𝑁, compute the sum of 𝑓 (𝐴, 𝑀) across all non-decreasing arrays of length 𝑁
containing elements in [0, 𝑁 − 1].

EXPLANATION:
Let’s fix the parameter 𝑀 and try to compute the sum of 𝑓 (𝐴, 𝑀) across all valid arrays 𝐴.

𝑓 (𝐴, 𝑀) is the number of distinct MEX-es attained; so to sum this up across all 𝐴, we can instead add up
the contributions of all the possible MEX-es.
That is, we do the following:

• Fix the MEX to be 𝐾.

• Count the number of valid arrays for which there exists a size 𝑀 subset with MEX equal to 𝐾.
• The MEX of an 𝑀-element set cannot exceed 𝑀, so we sum this up across all 0 ≤ 𝐾 ≤ 𝑀.

We now have a sub-problem: with 𝑀 and 𝐾 fixed, how many non-decreasing arrays 𝐴 with elements in
[0, 𝑁 − 1] contain some subset of size 𝑀 with MEX equal to 𝐾?

Here, note that dealing with non-decreasing arrays is equivalent to dealing with frequencies of elements.
Let 𝑓𝑥 denote the frequency of 𝑥.
Thinking in terms of frequencies is helpful here because it gives us quite straightforward conditions to
work with:

1 of 3 10/8/25, 23:23
Subset MEX Counting Practice Coding Problem [Link]

• For the MEX to equal 𝐾, certainly we must have (at least) one copy each of 0, 1, 2, …, 𝐾 − 1.
So, we require 𝑓0 , 𝑓1 , …, 𝑓𝐾 − 1 to all be positive.
• This gives us 𝐾 elements that certainly must be chosen as part of the subset already.
(𝑀 − 𝐾) more elements must be chosen to complete the subset with MEX 𝐾; and importantly, none
of the chosen elements must equal 𝐾.
This is possible only when 𝑓𝐾 , since if there are any more copies of 𝐾 in the array then by the
pigeonhole principle we’ll be forced to take at least one copy of 𝐾 into the subset.

The two conditions above, dealing with 𝑓0 , …, 𝑓𝐾 , are together necessary and sufficient for an array to
have a size-𝑀 subset with MEX equal to 𝐾.

Let’s now try to count such arrays - which as noted above is equivalent to counting frequency arrays
[𝑓0 , …, 𝑓𝑁 − 1 ].
Note that we have the following conditions on 𝑓 :

1. 𝑓0 + 𝑓1 + … + 𝑓𝑁 − 1 = 𝑁.
2. 𝑓𝑖 ≥ 0 for all 𝑖.
3. 𝑓𝑖 ≥ 1 for all 0 ≤ 𝑖 < 𝐾.
4. 𝑓𝐾 ≤ 𝑁 − 𝑀.

The last condition is the hardest one to deal with here - if we just had a target sum and lower bounds on
each element, the number of solutions would be easily found with the help of stars and bars.
Luckily, we can still work around the upper bound, by instead converting it to a lower bound.

To do this, let’s first ignore the upper bound entirely, and just count the number of sequences that satisfy
the first three conditions.
This is simply the number of non-negative integer solutions to

𝑓0 + 𝑓1 + … + 𝑓𝑁 − 1 = 𝑁 − 𝐾

which stars and bars tells us is equal to ( 𝑁 − 𝑁

𝐾 + 𝑁 − 1 = 2𝑁 − 𝐾 − 1 .
−1 ) ( 𝑁−1 )
Next, we subtract out the ‘bad’ configurations, which are all of those that have 𝑓𝐾 ≥ 𝑁 − 𝑀 + 1.
This is now a lower bound, so we’re again able to apply stars-and-bars, with the resulting count being
( 𝑁 + 𝑀𝑁 −−1𝐾 − 2 ).
So, with 𝐾 and 𝑀 fixed, the value we’re looking for is simply

( 2𝑁𝑁−−𝐾1− 1 ) − ( 𝑁 + 𝑀𝑁 −−1𝐾 − 2 )

For a fixed value of 𝑀, we want to quickly compute the sum

𝑀
∑𝐾 = 0 (( 2𝑁𝑁−−𝐾1− 1 ) − ( 𝑁 + 𝑀 −𝐾 −2
𝑁−1 ))
To do this quickly, note that every term is of the form ( 𝑁𝑥− 1 ), and we’re interested in the sum of these for
some contiguous range of 𝑥.
This allows for the use of prefix sums.
Specifically, define 𝑃𝑖 = ∑0 ≤ 𝑥 ≤ 𝑖 ( 𝑁𝑥− 1 ), and then the answer for 𝑀 simply becomes

𝑃2𝑁 − 1 − 𝑃2𝑁 − 𝑀 − 2 − 𝑃𝑁 + 𝑀 − 2 + 𝑃𝑁 − 3

Binomial coefficients can be computed in constant time after precomputation of factorials and inverse

2 of 3 10/8/25, 23:23
Subset MEX Counting Practice Coding Problem [Link]

factorials, so the answer for each 𝑀 can be found in 𝒪 (1) time, leading to 𝒪 (𝑁) overall.

TIME COMPLEXITY:
𝒪 (𝑁) per testcase.

CODE:
▾ Editorialist's code (PyPy3)
mod = 998244353
fac = [1] + list(range(1, 500005))
for i in range(1, 500005): fac[i] = fac[i-1] * i % mod
inv = fac[:]
inv[-1] = pow(inv[-1], mod-2, mod)
for i in reversed(range(500004)): inv[i] = inv[i+1] * (i+1) % mod
def C(n, r):
if n < r or r < 0: return 0
return (fac[n] * inv[r] % mod * inv[n-r]) % mod

for _ in range(int(input())):
n = int(input())
pref = [0]*(2*n)
for i in range(2*n):
pref[i] = C(i, n-1)
if i > 0: pref[i] = (pref[i] + pref[i-1]) % mod

ans = [0]*n
for m in range(1, n+1):
ans[m-1] = pref[2*n-1] + (pref[n-3] if n-3 >= 0 else 0)
ans[m-1] -= (pref[2*n-m-2] if 2*n-m-2 >= 0 else 0) + pref[n+m-2]
ans[m-1] %= mod
print(*ans)

3 of 3 10/8/25, 23:23