ABC290 C - Max MEX の考え方を書きました。

問題理解

MEX という関数が定義され、それを最大化するための考察と、高速なアルゴリズムでの実装が要求される問題でした。

MEX を想像しやすくするため具体例を出すと、

  • スーパーマーケットで、0番レジ、1番レジ、2番レジ… がある。
  • お客さんがN人いて、それぞれ [2、0、2、3、2、1、9 …] 番のレジに並んでいる。
  • あなたは0番レジ付近にいる。最も近い、空いているレジはどこか。

について、レジ番号が戻り値となるような関数です。

ABC290 C問題 は、N人のお客さんからK人だけを選ぶとき、MEXの最大でいくつにできるかという問題でした。

考察

MEXの定義から、以下のような性質があることがわかります。

  • 各レジに並んでいる具体的な人数の情報は不要で、“空いているか否か” だけが影響する。
  • 1~100000までが埋まっていたとしても、0番が空いていたら MEXは 0 となる。
    • 0から順に見ていき、穴があればその値が MEX となる。

これらをもとに、今回の問題においてMEXの最大値はいくつになるかを考えます。

この問題では、お客さんのうちK人しか選ぶことができません。
同じレジ番号のお客さんを2回以上選ぶ意味はなく、むしろ他のレジを埋める機会を失ってしまいます。
よって、レジ番号が重複しないように選んでいくことにします。

K人をうまく選択することで 0 ~ K-1 までを埋められるなら、MEX は K です。
例えば、 {0,1,2,3,4,5} から K=3 選んでMEXを最大化するならば、{0,1,2} を選び、MEX = 3 が最大値です。

アルゴリズム案

0~K-1 まで昇順に、それが A = {2, 0, 2, 3, 2, 1, 9 …} に含まれるか を判定していけば、MEXを求めることができます。
途中で含まれていないものが見つかればその値が答えですし、見つからなければ(すべて含まれているなら) K が答えです。

# イメージ

def solve():
    for i in range(K): # 0~K-1 を昇順ループ
        if i not in A:
            return i
    return K

計算量の見積もりと改善

“0~K-1 までを昇順に試す” ので、この時点で計算量は少なくとも O(K) です。

“A = {2, 0, 2, 3, 2, 1, 9 …} に含まれるか” の判定は、先頭から順に全要素について一致判定をしていく場合、 O(|A|) = O(N) です。
これでは、K回のループの中でN回のループをするため O(KN) となり、実行時間制限を超過してしまいます。(最大ケースで約10^14回のループ処理が発生することになりますが、この処理は2秒を超えてしまします。)
単に「含まれるか」の判定は、set(集合)や、辞書/Hash/Dict/Map と呼ばれる構造を使うと高速で、今回は実質 O(1) でできると考えてよいです。
従って、 O(K) で答えを求めることができます。
(入力受け取り部分などを考慮すると O(N)

実装

def solve(N, K, A):
    a = set(A)

    for i in range(0, K): # 0~K-1 を昇順ループ
        if i not in a:
            return i

    return K

回答例