はじめに
競技プログラミングをしていると、多重集合が存在して
1. 要素を追加したり
2. 最大値や最小値を求めたり
3. 集合内の要素を削除する
といったクエリをO(logN)で行わなければいけない場面にしばしば直面します。
その際、公式解説に「C++の標準ライブラリにあるMultisetを使って高速に～」といった記述を見ます。
しかしながら、Pythonの標準ライブラリにはMultisetは搭載されていません。
このような場合、通常Pythonでは標準ライブラリにあるheapqをつかって対処すること多いです。
この記事ではこのようなheapqを用いた疑似Multisetをクラス化したものについて記述します。
注意：この記事で実装する疑似MultisetではC++のようにMultiset上での二分探索はできません。
コード
とりあえずコードです。heapqを用いて実装します。
実装に関する詳しい説明は省略します。
#####疑似マルチセット QuasiMultiset#####
from heapq import heappop,heappush
from collections import defaultdict
class QuasiMultiset:
def __init__(self,deleteFlag=True):
'''
deleteFlag: 要素数が0となったkeyを辞書から削除するか
'''
self.len = 0
self.f = defaultdict(int)
self.pos = []
self.neg = []
self.deleteFlag = deleteFlag
def _delete(self,key):
'''
要素数が0となったkeyを辞書から削除
内部呼び出し用
'''
if self.deleteFlag and self.f[key] == 0:
del self.f[key]
def add(self,x,key=None):
'''
x:追加する要素の値
key:キー
'''
if key is None:
key = x
self.len += 1
heappush(self.pos,(x,key))
heappush(self.neg,(-x,key))
self.f[key] += 1
def popMax(self,returnKey=False):
'''
最大値とそのkey(optional)をpopする
'''
while True:
x,key = heappop(self.neg)
if self.f[key]:
self.f[key] -= 1
self.len -= 1
self._delete(key)
if returnKey:
return -x,key
else:
return x
def popMin(self,returnKey=False):
'''
最小値とそのkey(optional)をpopする
'''
while True:
x,key = heappop(self.pos)
if self.f[key]:
self.f[key] -= 1
self.len -= 1
self._delete(key)
if returnKey:
return x,key
else:
return x
def seachMax(self,returnKey=False):
'''
最大値とそのkey(optional)をreturnする
'''
while True:
x,key = self.neg[0]
if self.f[key]:
if returnKey:
return -x,key
else:
return -x
else:
heappop(self.neg)
def seachMin(self,returnKey=False):
'''
最小値とそのkey(optional)をreturnする
'''
while True:
x,key = self.pos[0]
if self.f[key]:
if returnKey:
return x,key
else:
return x
else:
heappop(self.pos)
def deleteKey(self,key):
'''
key:削除するkey
'''
self.f[key] -= 1
self._delete(key)
self.len -= 1
def exist(self,key):
if self.f.get(key) is None:
return False
else:
return True
なにができるのか
上述した疑似multisetで何ができるのかを説明します。
0. インスタンスの作成
multiset = QuasiMultiset()
1. addクエリ
疑似multisetに要素xを追加します。
multiset.add(x)
2. popMaxクエリ
疑似multiset内の最大値をpopします。
popされた最大値はmultisetから削除されます。
maxval = multiset.popMax()
3. popMinクエリ
疑似multiset内の最小値をpopします。
popされた最小値はmultisetから削除されます。
minval = multiset.popMin()
4. seachMaxクエリ
疑似multiset内の最大値を返します。
maxval = multiset.searchMax()
5. seachMinクエリ
疑似multiset内の最小値を返します。
minval = multiset.searchMin()
6. deleteクエリ
疑似multiset内の特定の値xを削除します。
multiset.delete(x)
7. existクエリ
疑似multiset内に特定の値xが存在するかbool型で返します。
b = multiset.exist(x)
計算時間
heapqを用いるので、すべてのクエリに対してO(logN)またはO(1)で行えます。(Nはmultiset内の要素数)
注意点
- 無効なクエリ（例：multisetが空なのにpopMaxをする）を行うとバグります。
- そこそこ遅いです。
- 冒頭にも述べましたが、c++などのmultisetにある二分探索ができないです。
使用例(ネタバレ注意)
ARC073 E. Ball Coloring
Code Festival 2018 Final F - Dinner Planning :メンバ変数lenはmultiset内の要素数を表します。