Official

B - Mex on Blackboard Editorial by evima


[1] Fixing the largest integer to write

Let \(x\) be the largest integer you write. You can write \(x\) only if the blackboard has \(0\) through \(x-1\). Here, if the initial blackboard does not have \(x-1\), you need to write \(x-1\). By repeating this recursive argument, we can see that you need to perform the following number of operations to write \(x\) for the first time:

  • the number of integers not on the initial blackboard + 1 (in which you write \(x\)).

[2] Counting

Let us try all possible values of \(x\).

After writing \(x\) for the first time, you can write any integer between \(0\) and \(x\). Additionally, from the above argument, there is a unique way to perform the minimum number of operations before writing \(x\) for the first time.

From these facts, the problem can be reduced to the following one, where \(M\) is the minimum number of operations needed to write \(x\) for the first time.


Let us write an integer between \(0\) and \(x\) exactly \(K-M\). How many multisets can be obtained in this way?


This is a famous problem of combination with repetition. There is a one-to-one correspondence between the set of such multisets and the ways to arrange \(K-M\) balls and \(x\) bars in a row, so there are \(\binom{K-M+x}{x}\) such multisets. This value can be computed in \(\mathrm{O}(1)\) time after pre-computing the factorials.

The possible values of \(x\) are \(0\) through \(N+K-1\), so the problem is solved in \(\mathrm{O}(N+K)\) time.


[3] Sample implementation

Python:

MOD = 998244353

table_len = 5 * 10**5

fac = [1, 1]
inv = [1, 1]
finv = [1, 1]
for i in range(2, table_len):
    fac.append(fac[-1] * i % MOD)
    inv.append(-inv[MOD % i] * (MOD // i) % MOD)
    finv.append(finv[-1] * inv[-1] % MOD)


def binom(n, r):
    return fac[n] * finv[n - r] % MOD * finv[r] % MOD


n, k = map(int, input().split())
a = list(map(int, input().split()))
exi = [0] * 400001
for v in a:
    exi[v] = 1
need, res = 1, 0
for i in range(n + k):
    if need > k:
        break
    rest = k - need
    res += binom(rest + i, i)
    res %= MOD
    if not exi[i]:
        need += 1
print(res)

posted:
last update: