priority_queue.py

# A min heap implementation with array
# Not tested yet
class PriorityQueue(object):

  # offset the index by one
  # this way, indexing are easier
  def __init__(self):
    self.arr = [None] 
  
  def __str__(self):
    return self.show()

  def show(self, joiner = ' '):
    return joiner.join([str(val) for val in self.arr])
  
  def peek(self):
    return self.__get(1) if not self.is_empty else None
  
  def push(self, val):
    self.arr.append(val)
    self.__swim_up(self.size)
  
  # Order is important
  # 1. swap first and last node
  # 2. remove the last (which previously was the first)
  # 3. swim the first down
  def pop(self):
    if self.is_empty:
      return None
    self.__swap(1, self.size)
    min_node = self.__rm(self.size)
    self.__swim_down(1)
    return min_node
  
  def __swim_up(self, i: int):
    parenti = self.__parent(i)
    if self.__has_parent(i) and self.__prior(i, parenti) == i:
      self.__swap(i, parenti)
      self.__swim_up(parenti)
  
  def __swim_down(self, i: int):
    if not self.__has_child(i):
      return
    # at this point, i has at least one child
    # that is, at least one of the children is not None
    lchildi, rchildi = self.__left_child(i), self.__right_child(i)
    p = self.__prior(i, self.__prior(lchildi, rchildi))
    # pick the child with the higher priority
    if p == lchildi:
      self.__swap(i, lchildi)
      self.__swim_down(lchildi)
    elif p == rchildi:
      self.__swap(i, rchildi)
      self.__swim_down(rchildi)
  
  def __get(self, index: int):
    if index >= self.size:
      return None
    return self.arr[index]
  def __set(self, index: int, val):
    while index >= self.size:
      self.arr.append(None)
    self.arr[index] = val
  def __rm(self, index: int):
    if index >= self.size:
      return None
    return self.arr.remove(index)
  def __swap(self, i: int, j: int):
    vali, valj = self.get(i), self.get(j)
    self.__set(i, valj)
    self.__set(j, vali)
  # Return the index with smaller value
  # assuming i and j are not both None
  # if one of them is none, simply return the other index
  def __prior(self, i: int, j: int):
    vali = self.get(i)
    if not vali:
      return j
    valj = self.get(j)
    if not valj:
      return i
    return i if vali < valj else j
  def __has_parent(self, i: int) -> bool:
    return i > 1
  def __has_child(self, i: int) -> bool:
    return self.__get(PriorityQueue.__left_child(i)) != None\
     or self.__get(PriorityQueue.__right_child(i)) != None
  @property
  def is_empty(self):
    return self.size == 0
  @property
  def size(self):
    return len(self.arr) - 1
  @staticmethod
  def __left_child(i: int) -> int:
    return 2 * i
  @staticmethod
  def __right_child(i: int) -> int:
    return PriorityQueue.__left_child(i) + 1
  @staticmethod
  def __parent(i: int) -> int:
    return i // 2