84 lines
2.7 KiB
Nim
84 lines
2.7 KiB
Nim
## 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()
|