両方向で見られる優先度付きキュー

最大値と最小値を両方O(logN)で見たい、ワガママな人のための優先度付きキュー。

要素の追加はO(logN)、要素の存在確認はO(1)
ただし、heapqを2つ使っているので使うメモリも2倍、計算量も2倍。

import heapq

class HeapBiq:
  def __init__(self,A=[]):
    self.low = A
    self.hi = [-1*a for a in A]
    heapq.heapify(self.low)
    heapq.heapify(self.hi)
    self.dic = {}
    for a in A:
      if a in self.dic:
        self.dic[a]+=1
      else:
        self.dic[a] = 1 
    
  def hpush(self,x):
    heapq.heappush(self.low,x)
    heapq.heappush(self.hi,-x)
    if x in self.dic:
      self.dic[x] += 1
    else:
      self.dic[x] = 1
    return
    
  def hpopmin(self):
    self.l = "error"
    while self.getlen() > 0:
      self.l = heapq.heappop(self.low)
      if self.l in self.dic:
        if self.dic[self.l] == 1:
          del self.dic[self.l]
        else:
          self.dic[self.l] -= 1
        break
    return self.l
    
  def hpopmax(self):
    self.h = "error"
    while self.getlen() > 0:
      self.h = -1*heapq.heappop(self.hi)
      if self.h in self.dic:
        if self.dic[self.h] == 1:
          del self.dic[self.h]
        else:
          self.dic[self.h] -= 1
        break
    return self.h
  
  def getlen(self):
    g = len(self.dic)
    return g

  def remove(self,x):
    if x in self.dic:
      if self.dic[x] == 1:
        del self.dic[x]
      else:
        self.dic[x] -= 1
    #存在しないxを削除しようとしてもerrorは返らない。
    return

  def getmin(self):
    if self.getlen() == 0:
      return "error"
    while True:
      ret = self.low[0]
      if ret in self.dic:
        return ret
      else:
        heapq.heappop(self.low)
    return
  
  def getmax(self):
    if self.getlen() == 0:
      return "error"
    while True:
      ret = -self.hi[0]
      if ret in self.dic:
        return ret
      else:
        heapq.heappop(self.hi)
    return

  def look(self):
    return self.dic

  def is_exist(self,x):
    if x in self.dic:
      return True
    else:
      return False


#使い方

A = [2,3,4,3,5,6,4]

hb = HeapBiq(A)
print(hb.hpopmin()) #2
print(hb.hpopmax()) #6
print(hb.look()) #{3: 2, 4: 2, 5: 1}

hb.remove(5)
print(hb.hpopmin()) #3
print(hb.hpopmax()) #4
hb.hpush(8)

print(hb.look()) #{3: 1, 4: 1, 8: 1}

print(hb.is_exist(9)) #False
print(hb.is_exist(4)) #True

print(hb.getmin()) #3
print(hb.getmax()) #8
print(hb.hpopmin()) #3
print(hb.hpopmax()) #8
print(hb.hpopmin()) #4
print(hb.look()) #{}
print(hb.hpopmax()) #error


参考
tsubo.hatenablog.jp