最大・最小・指定キー削除をヒープで実現する話

この記事でやること

Python3のheapqを使い、以下の操作ができるデータ構造を実装します：

1. push(x): xを追加
2. erase(x): xを削除
3. get_min(): 最小値を取得
4. get_max(): 最大値を取得

ソースコードはご自由にお使いください。

前提知識：（最小）ヒープとは

以下の操作ができるデータ構造です：

1. push(x): xを追加
2. pop(): 最小値を削除
3. get(): 最小値を取得

仕組みが知りたい方はwikipediaを参照してください。

（このwikiによると、指定キー削除は二分ヒープのノードを指すポインタを別に管理することでできるようなので、ヒープを一から自作するなら参考にするとよいと思います。）

発想

push(x), get_min(), get_max() 

最小ヒープと最大ヒープを両方持っておくだけで実現できます。heapqは最小ヒープなので、最大ヒープとして扱いたければデータを-1倍してpushすればいいでしょう。

erase(x)

遅延評価というテクニックを使って実現します。

実際にヒープから削除するのではなく、「xは削除済みである。」という情報だけ記憶しておきます。そして、後のget_min(), max_min()でもしxが取得されたなら、そのタイミングでヒープからpopします。

実装では、連想配列(defaultdict)で各要素の個数を管理しておき、個数が0のものがgetされたらヒープからpopするようにしています。

ソースコード

（Python3.8.3）

from collections import defaultdict
from heapq import heappush, heappop

class MinMaxHeap:
    
    def __init__(self, array=[]):
        self.hq_min = [] # 最小ヒープ
        self.hq_max = [] # 最大ヒープ
        self.cnt = defaultdict(int) # cnt[x] := xの個数
        self.size = 0 # 要素数
        for a in array:
            self.push(a)
    
    def push(self, x):
        # xを追加
        heappush(self.hq_min, x)
        heappush(self.hq_max, -x)
        self.cnt[x] += 1
        self.size += 1
    
    def erase(self, x):
        # xを削除（複数あれば1つだけ、無ければ何もしない）
        if self.cnt[x]:
            self.cnt[x] -= 1
            self.size -= 1
    
    def get_min(self):
        # 最小値を取得（空ならNone）
        while self.hq_min and self.cnt[self.hq_min[0]] == 0:
            heappop(self.hq_min)
        if self.hq_min:
            return self.hq_min[0]
        return None
    
    def get_max(self):
        # 最大値を取得（空ならNone）
        while self.hq_max and self.cnt[-self.hq_max[0]] == 0:
            heappop(self.hq_max)
        if self.hq_max:
            return -self.hq_max[0]
        return None
    
    def __len__(self):
        # len()で要素数を取得
        return self.size
    
    def __repr__(self):
        # {各要素: 個数} の形で表示される。
        return self.cnt

使用例：

ABC253-C問題

Submission #33418423 - NOMURA Programming Contest 2022（AtCoder Beginner Contest 253）