Pythonでmultisetの代用ができるデータ構造

ABC170のE問題でmultisetを使う問題が出題されましたが、Pythonにはmultisetがありません（さらに順序付きの普通の集合もないですが...）。そこで以下の処理を行うことができるデータ構造、HeapDictを作りました（ソースコードも一番下にあります）。

〇O(logn)で要素の挿入
〇O(logn)で要素の削除
〇O(1)で要素の存在確認
〇O(1)で最小値の取得

これは単純にheapqの上位互換です。C++のmultisetやsetとかなり似ています。二分探索ができないだけです。

PythonのsetやdictはO(1)で要素の挿入および削除ができますが、順序付きでないため最小値の取得にはO(n)かかります。また、PythonのheapqはO(logn)で要素の挿入ができ、O(1)で最小値の取得ができますが、最小値以外の削除ができない？はずです。しかし、これらを組み合わせることでmultiset的な挙動をするデータ構造を作ることができます。

具体的な実装方法を述べます。
heapqとdictを用意します。

def __init__(self):
    self.h=[]
    self.d=dict()

要素の挿入はどちらもできるため、普通に挿入します。計算量はdictでO(1)、heapqでO(logn)となります。

def insert(self,x):
    heapq.heappush(self.h,x)
    if x not in self.d:
        self.d[x]=1
    else:
        self.d[x]+=1

要素の削除はheapqでは普通はできませんが、削除を遅延させることで解決します。このデータ構造で取得したいものは最小値であるため、最小値さえ矛盾していなければ要素の削除は遅延させてもかまわないということになります。遅延方法としては、まずdictで普通に要素の削除をします。次にheapqの方の最小値を確認します。この最小値がdictにすでに存在しないとなれば削除し、最小値がdictに存在するようになるまで削除し続けます。遅延削除の回数は多くても通常の削除の回数と同じとなります。

def erase(self,x):
    if x not in self.d or self.d[x]==0:
        print(x,"is not in HeapDict")
        exit()
    else:
        self.d[x]-=1

    while len(self.h)!=0:
        if self.d[self.h[0]]==0:
            heapq.heappop(self.h)
        else:
            break

要素の存在確認はdictの方を見ればO(1)で行えます。

def is_exist(self,x):
    if x in self.d and self.d[x]!=0:
        return True
    else:
        return False

最小値の取得は、最小値が矛盾していないという保証があるため、heapqの方を見ればO(1)で行えます。

def get_min(self):
    return self.h[0]

最後にこれらをまとめたものがこちらになります。

import heapq
from sys import exit
class HeapDict:
    def __init__(self):
        self.h=[]
        self.d=dict()

    def insert(self,x):
        heapq.heappush(self.h,x)
        if x not in self.d:
            self.d[x]=1
        else:
            self.d[x]+=1

    def erase(self,x):
        if x not in self.d or self.d[x]==0:
            print(x,"is not in HeapDict")
            exit()
        else:
            self.d[x]-=1

        while len(self.h)!=0:
            if self.d[self.h[0]]==0:
                heapq.heappop(self.h)
            else:
                break

    def is_exist(self,x):
        if x in self.d and self.d[x]!=0:
            return True
        else:
            return False

    def get_min(self):
        return self.h[0]

使い方としてはPython標準のsetっぽく、insert(値)としたり、erase(値)としたりすればよいです。