ごりちゃんがゆく

競プロをするゴリラの精進記録

ABC330-E Mex and Update を無思考で解けるライブラリを書いた

Python アルゴリズム スニペット

ごり〜

こんばんは
先日のABC330-E を無思考で解けるライブラリを書いてみたので、紹介させていただきます！

問題概要は、 配列Aと、クエリとして i と x が与えられるので、
クエリのたびに A[i] を x に変更し、配列AのMEXを出力してください
といったものです。
MEXとは、配列に含まれない最小の非負整数をさします。

早速ですが、使い方はこんな感じになります

mex = Mex(A)
for i,x in qs:
    i -= 1
    mex.remove(A[i])
    mex.add(x)
    A[i] = x
    print(mex.get())

シンプルですね！
最近 MEXという概念がよく出るので、書いてみました！

ちなみに書いてみた動機は、AllDirectionsさんのポストを見ていいなあ〜と思ったからでした

動作原理とか

(解法のネタバレを含みます)

この問題を解こうとしたときに、どんなデータ構造があれば解けるのかな？と考えるはずです。
集合に含まれる要素の個数は管理できるとしても、含まれない非負整数のうち最小のものはどうやって求めるのか？と考えると思います。
実は、MEXというのは集合の要素数がNであるとき、Nより大きい値は取らないという特徴があります。
となれば、想定される最大のN までの非負整数のうち、集合に含まれないものを全部管理してしまえ！と発想できれば、解けた！に大きく近づきます
ところが、Python で解くには set で集合を管理できはしても高速に最小値を取れない、あるいは、heapqで 最小値を高速に取得できたとしても変更に弱いという難点があり、もう一山あります
(例えば C++ では std::set があります、Pythonでは tatyamさんの SortedSet や sortedcontainers を使うのが定跡のようです)
今回、これらを使わずとも、(機能を落としているぶん) 早く動作するモジュールを書きました！
原理としては、内部に heapq を2本持ち、本来 heapqでは削除できないところ、擬似的な削除を実現しています。

コードは下記で公開しています！
ぜひ使ってみてください！
github.com

from heapq import heappop,heappush,heapify
class Mex:    
    def __init__(self, arr=[], MAX=10**6) -> None:
        self.MAX = MAX + 5
        self.hist = [0] * (self.MAX+1)
        for a in arr:
            if a > self.MAX: continue
            self.hist[a] += 1
        self.q = []
        self.d = []
        for i in range(self.MAX+1):
            if self.hist[i] == 0:
                heappush(self.q, i)
    def get(self):
        while self.q and self.d and self.q[0]==self.d[0]:
            heappop(self.q)
            heappop(self.d)
        return self.q[0] if self.q else None
    def add(self, x):
        if x > self.MAX: return
        self.hist[x] += 1
        if self.hist[x] == 1:
            heappush(self.d, x)
    def remove(self, x):
        if x > self.MAX: return
        self.hist[x] -= 1
        if self.hist[x] == 0:
            heappush(self.q, x)

集合に含まれないものを管理する、というのは、MEXならではのある種の典型の気もします
が、そんなことも考えず、ライブラリ化されてると嬉しいよね！という気持ちです
verify提出

告知

ところで、告知 (というほど大げさなものではない) ですが、
「ごりちゃん競プロライブラリ」を始動します(した)
今後、ごりちゃんが便利と思ったコードはこちらにあげていく (つもり) (予定) なので、要チェックしてみてください！！
では！　良き競プロライフを！