## Generalized min/max binary heap implementation
## 2023/01/24
type BinHeap*[T] = ref object
  items: seq[T]
  size: int
  puCmp: proc (lhs, rhs: T): bool {.noSideEffect.}
  pdCmp: proc (lhs, rhs: T): bool {.noSideEffect.}

func lt[T](lhs, rhs: T): bool = lhs < rhs
func gt[T](lhs, rhs: T): bool = rhs < lhs

## Find the smallest children of a parent node
## Knowing the parent node, we can its children with this formula:
##   If P = 2, then its children are 2P and 2P + 1 respectively.
proc minChild[T](self: BinHeap[T], parent: int): int =
  let (lc, rc) = (parent * 2, parent * 2 + 1)
  # If the computed right child of the parent is bigger than
  # the heap list, it doesn't exist.
  if (rc >= self.items.len) or (self.items[lc] < self.items[rc]):
    return lc
  return rc

## Swaps the inserted element with its parent until the element is
## bigger than its parent, towards the root
proc propagateUp[T](self: BinHeap[T]; parent, child: var int) =
  while parent > 0:
    if self.puCmp(self.items[child], self.items[parent]):
      swap(self.items[parent], self.items[child])
    child = parent
    parent = parent div 2

## Swaps the node with its smallest children until the node is bigger
## than its parent, towards the leaves
proc propagateDown[T](self: BinHeap[T], node: int) =
  var i = node
  while i * 2 < self.items.len:
    let mc = self.minChild(i)
    if self.pdCmp(self.items[i], self.items[mc]):
      swap(self.items[i], self.items[mc])
    i = mc

proc newBinHeap*[T](minh: bool): BinHeap[T] =
  new(result)
  result.items = newSeqOfCap[T](32)
  result.items.add(0)
  result.puCmp = lt
  result.pdCmp = gt
  if not minh:
    swap(result.puCmp, result.pdCmp)

## Adds a new element to the heap, keeping the heap property.
proc add*[T](self: BinHeap[T], e: T) =
  self.items.add(e)
  inc self.size
  var
    child = self.items.len - 1 # the inserted element
    parent = child div 2       # parent of the element
  self.propagateUp(parent, child)

## Pops the min or max element (the root, index 1 in a min heap)
## keeping the heap property.
## Whether the popped element is min or max depends if the heap is a
## minumum or maximum one.
proc popMinOrMax*[T](self: BinHeap[T]): T =
  result = self.items[1]
  self.items[1] = self.items[self.items.len - 1]
  discard self.items.pop()
  self.propagateDown(1)

## Builds a heap from the given list
proc heapify*[T](self: BinHeap[T], lst: openArray[T]) =
  self.items &= lst
  for i in countdown(self.items.len div 2, 1):
    self.propagateDown(i)
    
func len*(self: BinHeap): int = self.size

when isMainModule:
  var bh = newBinHeap[int](false)
  bh.heapify([9, 6, 5, 2, 3])
  bh.add(23)
  for e in bh.items:
    echo e
  echo "min element: ", bh.popMinOrMax()