Official

B - Mex on Blackboard Editorial by nok0


[1] 操作で書く最大値の決め打ち

操作で書く最大値を \(x\) と決め打ちます。このとき、\(x\) を書くためには、黒板に \(0\) から \(x-1\) が書かれていることが必要です。ここで、初期状態の黒板に \(x-1\) が書かれていない場合は、\(x-1\) を書く必要があります。このように再帰的な議論を繰り返すことで、

  • \(0\) から \(x-1\) のうち、初期状態の黒板に書かれていない数の個数 + \(1\) 回(\(x\) を書く操作)

の操作が \(x\) をはじめて書くために必要だとわかります。


[2] 数え上げ

\(x\) を全探索しましょう。

\(x\) をはじめて書いた後の操作では、\(0\) から \(x\) までの任意の数を書くことができます。また、\(x\) をはじめて書くまでの最小回数の操作は上の議論から一意に定まります。

以上の事実から、\(x\) をはじめて書くまでに必要な操作回数を \(M\) とすると、この問題は以下の問題に帰着されます。


\(0\) 以上 \(x\) 以下の数を書く操作をちょうど \(K-M\) 回行う。このとき、得られる多重集合は何通りか。


これは重複組み合わせの有名な問題で、 \(K-M\) 個のボールと \(x\) 個のしきりを横一列に並べる方法と一対一対応し、\(\binom{K-M+x}{x}\) と計算できます。この値は、階乗の前計算のもと \(\mathrm{O}(1)\) で求められます。

\(x\) の範囲としては \(0\) 以上 \(N+K\) 未満を調べればいいことから、この問題を \(\mathrm{O}(N+K)\) で解くことができました。


[3] 実装例

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: