6
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

More than 3 years have passed since last update.

PythonでMultisetっぽいことをする(heapqを用いた実装)

Last updated at Posted at 2021-08-31

#はじめに
競技プログラミングをしていると、多重集合が存在して
 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内の要素数)

#注意点

  1. 無効なクエリ(例:multisetが空なのにpopMaxをする)を行うとバグります。
  2. そこそこ遅いです。
  3. 冒頭にも述べましたが、c++などのmultisetにある二分探索ができないです。

#使用例(ネタバレ注意)
ARC073 E. Ball Coloring
Code Festival 2018 Final F - Dinner Planning :メンバ変数lenはmultiset内の要素数を表します。

6
3
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
6
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?