【Python】削除可能優先度付きキュー

競技プログラミング

機能

通常の優先度付きキューの機能
- 要素の追加：O(logN)
- 最小値（最大値）の出力または削除：O(logN)
に加えて、
- 特定の要素の削除：O(logN)
をすることができます。

使い方

  • キューのインスタンス作成
    • 最小値の取得をするキュー：q = DeletablePQ()（入れたいデータがある場合はq=DeletablePQ(data=arr)）
    • 最大値の取得をするキュー：q = DeletablePQ(greater=True)（入れたいデータがある場合はq=DeletablePQ(data=arr,greater=True)）
  • 要素xの追加
    • q.push(x)
  • 要素xの削除
    • q.remove(x)
  • 要素xをn個削除
    • q.nremove(x, num=n)
  • 最小値（最大値）を削除（返り値あり）
    • x = q.pop()
  • xの個数を出力
    • q.count(x)
  • 最小値（最大値）を出力（削除はせず、参照するだけ）
    • x = q.get_val()
  • 要素の総数を出力
    • len(q)
  • 他のDeletablePQのインスタンスとマージ（要素数の多いインスタンスをベースに行うとより速いです）
    • q.merge(other_q)

使用例

from heapq import heappush,heappop
 
class DeletablePQ:
  def __init__(self,data=None,greater=False):
    self.greater=greater
    self.count= {}
    self.que=[]
    if data:
        for x in data:
            heappush(self.que,x*(-1 if greater else 1))
            if self.count.get(x): self.count[x] += 1
            else: self.count[x] = 1
    self.len = len(self.que)
 
  def push(self,x):
    heappush(self.que,x*(-1 if self.greater else 1))
    if self.count.get(x): self.count[x] += 1
    else: self.count[x] = 1
    self.len += 1
    
  def remove(self,x):
    self.count[x] -= 1
    self.len -= 1
  
  def nremove(self,x,num=1):
    self.count[x] -= num
    self.len -= num
    if self.count[x] < 0:
      self.len -= self.count[x]
      self.count[x] = 0
    
  def pop(self):
    x,self.count[x]=-1,0
    while self.count[x]<1: x=heappop(self.que)*(-1 if self.greater else 1)
    self.len -= 1
    self.count[x] -= 1
    return x
  
  def count(self,x):
    return self.count[x] if self.count[x]>0 else 0
  
  def get_val(self):
    while self.count.get(self.que[0]*(-1 if self.greater else 1)) != None and self.count[self.que[0]*(-1 if self.greater else 1)]<1:
        heappop(self.que)
    return self.que[0]*(-1 if self.greater else 1)
 
  def merge(self, other):
        while len(other) > 0: self.push(other.pop())
  
  def __len__(self):
    return self.len