Improve the heapqueue module (#17034)

Improve documentation
Optimize toHeapQueue
Rename siftup and siftdown
Add tests for the heap property

lib/pure/collections/heapqueue.nim

@@ -8,31 +8,27 @@
 
 
 ## The `heapqueue` module implements a
-## `heap data structure<https://en.wikipedia.org/wiki/Heap_(data_structure)>`_
+## `binary heap data structure<https://en.wikipedia.org/wiki/Binary_heap>`_
-## that can be used as a
+## that can be used as a `priority queue<https://en.wikipedia.org/wiki/Priority_queue>`_.
-## `priority queue<https://en.wikipedia.org/wiki/Priority_queue>`_.
+## They are represented as arrays for which `a[k] <= a[2*k+1]` and `a[k] <= a[2*k+2]`
-## Heaps are arrays for which `a[k] <= a[2*k+1]` and `a[k] <= a[2*k+2]` for
+## for all indices `k` (counting elements from 0). The interesting property of a heap is that
-## all `k`, counting elements from 0. The interesting property of a heap is that
 ## `a[0]` is always its smallest element.
 ##
 ## Basic usage
 ## -----------
 ##
-
 runnableExamples:
-  var heap = initHeapQueue[int]()
+  var heap = [8, 2].toHeapQueue
-  heap.push(8)
-  heap.push(2)
   heap.push(5)
-  # The first element is the lowest element
+  # the first element is the lowest element
   assert heap[0] == 2
-  # Remove and return the lowest element
+  # remove and return the lowest element
   assert heap.pop() == 2
-  # The lowest element remaining is 5
+  # the lowest element remaining is 5
   assert heap[0] == 5
 
-## Usage with custom object
+## Usage with custom objects
-## ------------------------
+## -------------------------
 ## To use a `HeapQueue` with a custom object, the `<` operator must be
 ## implemented.
 
@@ -48,6 +44,7 @@ runnableExamples:
 
   assert jobs[0].priority == 1
 
+
 import std/private/since
 
 type HeapQueue*[T] = object
@@ -57,27 +54,33 @@ type HeapQueue*[T] = object
 proc initHeapQueue*[T](): HeapQueue[T] =
   ## Creates a new empty heap.
   ##
-  ## See also:
+  ## Heaps are initialized by default, so it is not necessary to call
+  ## this function explicitly.
+  ##
+  ## **See also:**
   ## * `toHeapQueue proc <#toHeapQueue,openArray[T]>`_
   discard
 
 proc len*[T](heap: HeapQueue[T]): int {.inline.} =
   ## Returns the number of elements of `heap`.
+  runnableExamples:
+    let heap = [9, 5, 8].toHeapQueue
+    assert heap.len == 3
+
   heap.data.len
 
 proc `[]`*[T](heap: HeapQueue[T], i: Natural): lent T {.inline.} =
   ## Accesses the i-th element of `heap`.
   heap.data[i]
 
-proc heapCmp[T](x, y: T): bool {.inline.} =
+proc heapCmp[T](x, y: T): bool {.inline.} = x < y
-  return (x < y)
 
-proc siftdown[T](heap: var HeapQueue[T], startpos, p: int) =
+proc siftup[T](heap: var HeapQueue[T], startpos, p: int) =
-  ## 'heap' is a heap at all indices >= startpos, except possibly for `pos`. `pos`
+  ## `heap` is a heap at all indices >= `startpos`, except possibly for `p`. `p`
   ## is the index of a leaf with a possibly out-of-order value. Restores the
   ## heap invariant.
   var pos = p
-  var newitem = heap[pos]
+  let newitem = heap[pos]
   # Follow the path to the root, moving parents down until finding a place
   # newitem fits.
   while pos > startpos:
@@ -90,7 +93,8 @@ proc siftdown[T](heap: var HeapQueue[T], startpos, p: int) =
       break
   heap.data[pos] = newitem
 
-proc siftup[T](heap: var HeapQueue[T], p: int) =
+proc siftdownToBottom[T](heap: var HeapQueue[T], p: int) =
+  # This is faster when the element should be close to the bottom.
   let endpos = len(heap)
   var pos = p
   let startpos = pos
@@ -109,48 +113,67 @@ proc siftup[T](heap: var HeapQueue[T], p: int) =
   # The leaf at pos is empty now. Put newitem there, and bubble it up
   # to its final resting place (by sifting its parents down).
   heap.data[pos] = newitem
-  siftdown(heap, startpos, pos)
+  siftup(heap, startpos, pos)
+
+proc siftdown[T](heap: var HeapQueue[T], p: int) =
+  let endpos = len(heap)
+  var pos = p
+  let newitem = heap[pos]
+  var childpos = 2 * pos + 1
+  while childpos < endpos:
+    let rightpos = childpos + 1
+    if rightpos < endpos and not heapCmp(heap[childpos], heap[rightpos]):
+      childpos = rightpos
+    if not heapCmp(heap[childpos], newitem):
+      break
+    heap.data[pos] = heap[childpos]
+    pos = childpos
+    childpos = 2 * pos + 1
+  heap.data[pos] = newitem
 
 proc push*[T](heap: var HeapQueue[T], item: sink T) =
-  ## Pushes `item` onto heap, maintaining the heap invariant.
+  ## Pushes `item` onto `heap`, maintaining the heap invariant.
   heap.data.add(item)
-  siftdown(heap, 0, len(heap)-1)
+  siftup(heap, 0, len(heap) - 1)
 
 proc toHeapQueue*[T](x: openArray[T]): HeapQueue[T] {.since: (1, 3).} =
   ## Creates a new HeapQueue that contains the elements of `x`.
   ##
-  ## See also:
+  ## **See also:**
   ## * `initHeapQueue proc <#initHeapQueue>`_
   runnableExamples:
-    var heap = toHeapQueue([9, 5, 8])
+    var heap = [9, 5, 8].toHeapQueue
     assert heap.pop() == 5
     assert heap[0] == 8
 
-  result = initHeapQueue[T]()
+  # see https://en.wikipedia.org/wiki/Binary_heap#Building_a_heap
-  for item in items(x):
+  result.data = @x
-    result.push(item)
+  for i in countdown(x.len div 2 - 1, 0):
+    siftdown(result, i)
 
 proc pop*[T](heap: var HeapQueue[T]): T =
   ## Pops and returns the smallest item from `heap`,
   ## maintaining the heap invariant.
   runnableExamples:
-    var heap = toHeapQueue([9, 5, 8])
+    var heap = [9, 5, 8].toHeapQueue
     assert heap.pop() == 5
+
   let lastelt = heap.data.pop()
   if heap.len > 0:
     result = heap[0]
     heap.data[0] = lastelt
-    siftup(heap, 0)
+    siftdownToBottom(heap, 0)
   else:
     result = lastelt
 
 proc find*[T](heap: HeapQueue[T], x: T): int {.since: (1, 3).} =
-  ## Linear scan to find index of item ``x`` or -1 if not found.
+  ## Linear scan to find the index of the item `x` or -1 if not found.
   runnableExamples:
-    var heap = toHeapQueue([9, 5, 8])
+    let heap = [9, 5, 8].toHeapQueue
     assert heap.find(5) == 0
     assert heap.find(9) == 1
     assert heap.find(777) == -1
+
   result = -1
   for i in 0 ..< heap.len:
     if heap[i] == x: return i
@@ -158,65 +181,69 @@ proc find*[T](heap: HeapQueue[T], x: T): int {.since: (1, 3).} =
 proc del*[T](heap: var HeapQueue[T], index: Natural) =
   ## Removes the element at `index` from `heap`, maintaining the heap invariant.
   runnableExamples:
-    var heap = toHeapQueue([9, 5, 8])
+    var heap = [9, 5, 8].toHeapQueue
     heap.del(1)
     assert heap[0] == 5
     assert heap[1] == 8
+
   swap(heap.data[^1], heap.data[index])
   let newLen = heap.len - 1
   heap.data.setLen(newLen)
   if index < newLen:
-    heap.siftup(index)
+    siftdownToBottom(heap, index)
 
 proc replace*[T](heap: var HeapQueue[T], item: sink T): T =
   ## Pops and returns the current smallest value, and add the new item.
-  ## This is more efficient than pop() followed by push(), and can be
+  ## This is more efficient than `pop()` followed by `push()`, and can be
   ## more appropriate when using a fixed-size heap. Note that the value
-  ## returned may be larger than item! That constrains reasonable uses of
+  ## returned may be larger than `item`! That constrains reasonable uses of
-  ## this routine unless written as part of a conditional replacement:
+  ## this routine unless written as part of a conditional replacement.
+  ##
+  ## **See also:**
+  ## * `pushpop proc <#pushpop,HeapQueue[T],sinkT>`_
   runnableExamples:
-    var heap = initHeapQueue[int]()
+    var heap = [5, 12].toHeapQueue
-    heap.push(5)
-    heap.push(12)
     assert heap.replace(6) == 5
     assert heap.len == 2
     assert heap[0] == 6
     assert heap.replace(4) == 6
+
   result = heap[0]
   heap.data[0] = item
-  siftup(heap, 0)
+  siftdown(heap, 0)
 
 proc pushpop*[T](heap: var HeapQueue[T], item: sink T): T =
-  ## Fast version of a push followed by a pop.
+  ## Fast version of a `push()` followed by a `pop()`.
+  ##
+  ## **See also:**
+  ## * `replace proc <#replace,HeapQueue[T],sinkT>`_
   runnableExamples:
-    var heap = initHeapQueue[int]()
+    var heap = [5, 12].toHeapQueue
-    heap.push(5)
-    heap.push(12)
     assert heap.pushpop(6) == 5
     assert heap.len == 2
     assert heap[0] == 6
     assert heap.pushpop(4) == 4
+
   result = item
   if heap.len > 0 and heapCmp(heap.data[0], result):
     swap(result, heap.data[0])
-    siftup(heap, 0)
+    siftdown(heap, 0)
 
 proc clear*[T](heap: var HeapQueue[T]) =
   ## Removes all elements from `heap`, making it empty.
   runnableExamples:
-    var heap = initHeapQueue[int]()
+    var heap = [9, 5, 8].toHeapQueue
-    heap.push(1)
     heap.clear()
     assert heap.len == 0
+
   heap.data.setLen(0)
 
 proc `$`*[T](heap: HeapQueue[T]): string =
   ## Turns a heap into its string representation.
   runnableExamples:
-    var heap = initHeapQueue[int]()
+    let heap = [1, 2].toHeapQueue
-    heap.push(1)
-    heap.push(2)
     assert $heap == "[1, 2]"
+
   result = "["
   for x in heap.data:
     if result.len > 1: result.add(", ")

tests/stdlib/theapqueue.nim

@@ -1,4 +1,4 @@
-import heapqueue
+import std/heapqueue
 
 
 proc toSortedSeq[T](h: HeapQueue[T]): seq[T] =
@@ -7,7 +7,18 @@ proc toSortedSeq[T](h: HeapQueue[T]): seq[T] =
   while tmp.len > 0:
     result.add(pop(tmp))
 
-block: # Simple sanity test
+proc heapProperty[T](h: HeapQueue[T]): bool =
+  for k in 0 .. h.len - 2: # the last element is always a leaf
+    let left = 2 * k + 1
+    if left < h.len and h[left] < h[k]:
+      return false
+    let right = left + 1
+    if right < h.len and h[right] < h[k]:
+      return false
+  true
+
+template main() =
+  block: # simple sanity test
     var heap = initHeapQueue[int]()
     let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
     for item in data:
@@ -16,7 +27,7 @@ block: # Simple sanity test
     doAssert(heap[0] == 0)
     doAssert(heap.toSortedSeq == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
 
-block: # Test del
+  block: # test del
     var heap = initHeapQueue[int]()
     let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
     for item in data: push(heap, item)
@@ -38,7 +49,7 @@ block: # Test del
 
     doAssert(heap.find(2) == -1)
 
-block: # Test del last
+  block: # test del last
     var heap = initHeapQueue[int]()
     let data = [1, 2, 3]
     for item in data: push(heap, item)
@@ -51,3 +62,41 @@ block: # Test del last
 
     heap.del(0)
     doAssert(heap.toSortedSeq == @[])
+
+  block: # testing the heap proeprty
+    var heap = [1, 4, 2, 5].toHeapQueue
+    doAssert heapProperty(heap)
+
+    heap.push(42)
+    doAssert heapProperty(heap)
+    heap.push(0)
+    doAssert heapProperty(heap)
+    heap.push(3)
+    doAssert heapProperty(heap)
+    heap.push(3)
+    doAssert heapProperty(heap)
+
+    # [0, 3, 1, 4, 42, 2, 3, 5]
+
+    discard heap.pop()
+    doAssert heapProperty(heap)
+    discard heap.pop()
+    doAssert heapProperty(heap)
+
+    heap.del(2)
+    doAssert heapProperty(heap)
+
+    # [2, 3, 5, 4, 42]
+
+    discard heap.replace(12)
+    doAssert heapProperty(heap)
+    discard heap.replace(1)
+    doAssert heapProperty(heap)
+
+    discard heap.pushpop(2)
+    doAssert heapProperty(heap)
+    discard heap.pushpop(0)
+    doAssert heapProperty(heap)
+
+static: main()
+main()