ABC330-E　heapq(優先度付きキュー(Priority queue))を利用した解法

Atcoder Beginner Contest 330のE問題「Mex and Update」について、Pythonにおけるheapqアルゴリズム（優先度付きキューアルゴリズム）を用いた解法に関する解説が見当たらなかったので、茶色コーダーの分際ながらネットの海に放浪させておきます。

問題

AtCoder Beginner Contest 330 E - Mex and Update

E - Mex and Update

長さNの数列A（A1,A2,・・・,An）がある。
以下のようなQ個のクエリに対応せよ。
与えられるクエリはそれぞれ以下の形式：「i, x」
①Aのi番目の値をxに変更する。
②AのMexを出力する。（Mex：Aに含まれない最小の非負整数）

考え方

①の値の変更については、特に重い操作ではないので問題ありません。

問題は、②のMexを求める作業です。
Mexの値をクエリのたびに純粋に求めようとする場合、非負整数を小さい順に、その値がAの中にあるかどうかを調べていくことになります。こうしてしまうと、最悪で0からNまでのN回をクエリQ個それぞれで調べることになり、計算量(NQ=4*10^10)でTLEになります。

ということで、高速にMexを求めるために何かしらの情報を別に持っておきたくなりました。
ここで、Aの中に”ない”値を何かしらの配列（ここでは配列Bと名付けておきます）で持っておくことを考えます。そうすることで、配列Bの中にある最小の値をMexとして採用できそうです。また、クエリごとのAへの操作と配列Bを対応させることについても、操作によってAの中の個数が0個になった値を配列Bに追加し、操作によって今までAに1つもなかった値がAに追加された場合はその値を配列Bから削除することで、難なくできそうです。

結局どういう配列が欲しいかを整理すると、上記に出てきた配列の
・Aの現在の状態（入力されたAを実際に操作していく）
・最小の値を素早く返してくれる配列（上記の配列B）
を持っておきます。
またそれに加えて、Aから消された結果Aの中の個数が0個になるタイミングや、Aの中の個数が0個かどうかを素早く把握したいので、
・Aにそれぞれの値がいくつ入っているのかを数えておく配列
も持っておきます。

解答

コンテスト中に提出したものではありませんが、体裁を整えただけで中身は大して変わっていません。

Submission #48184596 - TOYOTA SYSTEMS Programming Contest 2023(AtCoder Beginner Contest 330)

from collections import defaultdict
import heapq

n,q = map(int, input().split())
a = list(map(int, input().split()))
s = defaultdict(int)
for i in a:
    s[i] += 1

t = []
for i in range(n+5):
    if s[i] == 0:
        heapq.heappush(t,i)
        
for _ in range(q):
    ind,num = map(int, input().split())
    ind -= 1

    s[a[ind]] -= 1
    if s[a[ind]] == 0:
        heapq.heappush(t, a[ind])

    s[num] += 1
    a[ind] = num
    while 1:
        u = heapq.heappop(t)
        if s[u] == 0:
            print(u)
            heapq.heappush(t,u)
            break

スマートなコードではないですが、要点だけ解説を付しておきます。

n,q = map(int, input().split())
a = list(map(int, input().split()))
s = defaultdict(int)
for i in a:
    s[i] += 1

sはint型のdefaultdictです。詳しくはご自身で調べていただければ幸いなのですが、簡単に言うとここでは「値：その個数（int）」の辞書を作れるようなデータ構造です。

ここではaにあるそれぞれの値の数を持っておきます。
このあとのクエリで順次変更していきます。

t = []
for i in range(n+5):
    if s[i] == 0:
        heapq.heappush(t,i)

tをheapq（優先度付きキュー）としておきます。
tはいわばMex候補の控え室です。0からn+5までの値のうち、aに入ることができていないあぶれ者を集めておきます。（厳密にはn+5ではなくnまででよかったのですが、コンテスト中はそのあたりちゃんと検証するのがめんどくさかったので…ちゃんとした検討は公式解説を参照ください。）
0からn+4までそれぞれiとして、もしその値がaの中になければ(すなわちS[i] == 0なら)、tに追加します。
これも後のクエリで順次変更していきます。

for _ in range(q):
    ind,num = map(int, input().split())
    ind -= 1

indはインデックスですので、始点0に変更しておきます。

    s[a[ind]] -= 1
    if s[a[ind]] == 0:
        heapq.heappush(t, a[ind])

sのa[ind]の値(このあと変更されて消える値)を-1します。
これによってs[a[ind]]==0になった場合（すなわちaからその値が1つもなくなった場合）、a[ind]をあぶれ者の待機場所tに追加します。

    s[num] += 1
    a[ind] = num

変更する前準備が整ったので、必要な値を変更後の値に修正していきます。
s[変更先の値]を1つ増やし、aも書き換えます。

    while 1:
        u = heapq.heappop(t)
        if s[u] == 0:
            print(u)
            heapq.heappush(t,u)
            break

Mexの値を確認します。

tに入っている候補を小さい順に確認していきます。ここまで、aから消えた値をheapqに追加する操作は行っていますが、aに再び追加された値は追えてませんので、tの最小値がそのままMexというわけではありません。
uとして、tの最小値を取り出してみます。
uがaの中にある値なら特に何もせずスルー。そのまま控え室から出て行ってもらいます。
uが本当にaの中に1つもない値なら、それを出力して、またtに戻して、このクエリは終わり。
（ここはもっとスマートに書けそうですが、とりあえずACできたのでヨシ！）
これで、すべての変更されるべき配列を変更し終わっています。

解説は以上です。
解説の皆々様、C++でできることをPythonでどうやるかも教えてくれ～～！！