heapqって？？？

最近少しatcoderにはまっていまして、書きなれたpythonで参加しています。 しかし、先日のABCで全くLTEが解消できない問題に出くわしてしまいました。そこでコンテスト後にtwitterで「ABC253 C問題」で検索をかけていると有識者らしき人たちがみなheapq, heapqと唱えているじゃありませんか。 そこで今回はheapqについて調べたものをまとめてみたいと思います。

heapqとは

このサイトによるとheapqとはpythonで用意された「優先度付きキュー（？）」の標準ライブラリらしいです。 まず、「優先度付きキュー」から。上記のサイトによると

優先度付きキュー (Priority queue) はデータ型の一つで、具体的には

• 最小値(最大値)をO(log N)で取り出す
• 要素をO(log N)で挿入する

ことができます。

らしいです。要するに「最小値を取り出す」「要素を挿入する」という操作が通常のリストに比べて早くできるということみたいですね。(通常のリストはO(N))

今回の問題（ABC253 C問題）

問題文

今回の問題は以下の通りです。

整数の多重集合 S があります。はじめ S は空です。 Q 個のクエリが与えられるので順に処理してください。 クエリは次の 3 種類のいずれかです。

• 1 x：Sにxを1個追加する
• 2 x c：Sからxをmin(c, (Sに含まれるxの個数))個削除する
• 3：(Sの最大値)-(Sの最小値)を出力する。このクエリを処理するとき、Sが空でないことが保証される

制約については元のサイトを参照してみてください。

解答を眺める

import heapq
from collections import defaultdict

mx=[]
mn=[]
cnt=defaultdict(int)

q=int(input())
for _ in range(q):
  query=list(map(int,input().split()))
  if query[0]==1:
    x=query[1]
    cnt[x]+=1
    heapq.heappush(mx,-x)
    heapq.heappush(mn,x)

  if query[0]==2:
    x,c=query[1:]
    cnt[x]=max(0,cnt[x]-c)

  if query[0]==3:
    while cnt[-mx[0]]==0:
      heapq.heappop(mx)
    while cnt[mn[0]]==0:
      heapq.heappop(mn)
    print(-mx[0]-mn[0])

heapqならではのメソッドが出てきたのでこれらについて確認します。

• .heappush(heapq_obj,x)：heapq_objにxを挿入します。
• .heapqpop(heapq_obj)：heapq_objから最小値を取り出します。
• .heapify(list)：リストを優先度付きキューへ変換します。

ここで注意が必要なのが優先度付きキューとなってもデータ型はpythonデフォルトのリストというところです。 優先度付きキューの原理については後日追記したいと思います。

heapqpopのテスト結果

import heapq

a = [3, 5, 1, 2, 0]
heapq.heapify(a)
for i in range(5):
    print(f"pop:{heapq.heappop(a)}")
    print(a)

# pop: 0
# [1, 2, 5, 3]
# pop: 1
# [2, 3, 5]
# pop: 2
# [3, 5]
# pop: 3
# [5]
# pop: 5
# []

確かに最小値をうまく出力できている（なんで...？）