mirror of
https://github.com/nim-lang/Nim.git
synced 2025-12-28 17:04:41 +00:00
Improve the heapqueue module (#17034)
Improve documentation Optimize toHeapQueue Rename siftup and siftdown Add tests for the heap property
This commit is contained in:
konsumlamm
@@ -8,31 +8,27 @@
|
||||
|
||||
|
||||
## The `heapqueue` module implements a
|
||||
## `heap data structure<https://en.wikipedia.org/wiki/Heap_(data_structure)>`_
|
||||
## that can be used as a
|
||||
## `priority queue<https://en.wikipedia.org/wiki/Priority_queue>`_.
|
||||
## Heaps are arrays for which `a[k] <= a[2*k+1]` and `a[k] <= a[2*k+2]` for
|
||||
## all `k`, counting elements from 0. The interesting property of a heap is that
|
||||
## `binary heap data structure<https://en.wikipedia.org/wiki/Binary_heap>`_
|
||||
## that can be used as a `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]`
|
||||
## for all indices `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]()
|
||||
heap.push(8)
|
||||
heap.push(2)
|
||||
var heap = [8, 2].toHeapQueue
|
||||
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.} =
|
||||
return (x < y)
|
||||
proc heapCmp[T](x, y: T): bool {.inline.} = x < y
|
||||
|
||||
proc siftdown[T](heap: var HeapQueue[T], startpos, p: int) =
|
||||
## 'heap' is a heap at all indices >= startpos, except possibly for `pos`. `pos`
|
||||
proc siftup[T](heap: var HeapQueue[T], startpos, p: int) =
|
||||
## `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,13 +93,14 @@ 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
|
||||
let newitem = heap[pos]
|
||||
# Bubble up the smaller child until hitting a leaf.
|
||||
var childpos = 2*pos + 1 # leftmost child position
|
||||
var childpos = 2 * pos + 1 # leftmost child position
|
||||
while childpos < endpos:
|
||||
# Set childpos to index of smaller child.
|
||||
let rightpos = childpos + 1
|
||||
@@ -105,52 +109,71 @@ proc siftup[T](heap: var HeapQueue[T], p: int) =
|
||||
# Move the smaller child up.
|
||||
heap.data[pos] = heap[childpos]
|
||||
pos = childpos
|
||||
childpos = 2*pos + 1
|
||||
# The leaf at pos is empty now. Put newitem there, and bubble it up
|
||||
childpos = 2 * pos + 1
|
||||
# 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]()
|
||||
for item in items(x):
|
||||
result.push(item)
|
||||
# see https://en.wikipedia.org/wiki/Binary_heap#Building_a_heap
|
||||
result.data = @x
|
||||
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
|
||||
## this routine unless written as part of a conditional replacement:
|
||||
## returned may be larger than `item`! That constrains reasonable uses of
|
||||
## this routine unless written as part of a conditional replacement.
|
||||
##
|
||||
## **See also:**
|
||||
## * `pushpop proc <#pushpop,HeapQueue[T],sinkT>`_
|
||||
runnableExamples:
|
||||
var heap = initHeapQueue[int]()
|
||||
heap.push(5)
|
||||
heap.push(12)
|
||||
var heap = [5, 12].toHeapQueue
|
||||
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]()
|
||||
heap.push(5)
|
||||
heap.push(12)
|
||||
var heap = [5, 12].toHeapQueue
|
||||
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]()
|
||||
heap.push(1)
|
||||
var heap = [9, 5, 8].toHeapQueue
|
||||
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]()
|
||||
heap.push(1)
|
||||
heap.push(2)
|
||||
let heap = [1, 2].toHeapQueue
|
||||
assert $heap == "[1, 2]"
|
||||
|
||||
result = "["
|
||||
for x in heap.data:
|
||||
if result.len > 1: result.add(", ")
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import heapqueue
|
||||
import std/heapqueue
|
||||
|
||||
|
||||
proc toSortedSeq[T](h: HeapQueue[T]): seq[T] =
|
||||
@@ -7,47 +7,96 @@ proc toSortedSeq[T](h: HeapQueue[T]): seq[T] =
|
||||
while tmp.len > 0:
|
||||
result.add(pop(tmp))
|
||||
|
||||
block: # Simple sanity test
|
||||
var heap = initHeapQueue[int]()
|
||||
let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
|
||||
for item in data:
|
||||
push(heap, item)
|
||||
doAssert(heap == data.toHeapQueue)
|
||||
doAssert(heap[0] == 0)
|
||||
doAssert(heap.toSortedSeq == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
|
||||
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
|
||||
|
||||
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)
|
||||
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:
|
||||
push(heap, item)
|
||||
doAssert(heap == data.toHeapQueue)
|
||||
doAssert(heap[0] == 0)
|
||||
doAssert(heap.toSortedSeq == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
|
||||
|
||||
heap.del(0)
|
||||
doAssert(heap[0] == 1)
|
||||
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)
|
||||
|
||||
heap.del(heap.find(7))
|
||||
doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 5, 6, 8, 9])
|
||||
heap.del(0)
|
||||
doAssert(heap[0] == 1)
|
||||
|
||||
heap.del(heap.find(5))
|
||||
doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 6, 8, 9])
|
||||
heap.del(heap.find(7))
|
||||
doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 5, 6, 8, 9])
|
||||
|
||||
heap.del(heap.find(6))
|
||||
doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 8, 9])
|
||||
heap.del(heap.find(5))
|
||||
doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 6, 8, 9])
|
||||
|
||||
heap.del(heap.find(2))
|
||||
doAssert(heap.toSortedSeq == @[1, 3, 4, 8, 9])
|
||||
heap.del(heap.find(6))
|
||||
doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 8, 9])
|
||||
|
||||
doAssert(heap.find(2) == -1)
|
||||
heap.del(heap.find(2))
|
||||
doAssert(heap.toSortedSeq == @[1, 3, 4, 8, 9])
|
||||
|
||||
block: # Test del last
|
||||
var heap = initHeapQueue[int]()
|
||||
let data = [1, 2, 3]
|
||||
for item in data: push(heap, item)
|
||||
doAssert(heap.find(2) == -1)
|
||||
|
||||
heap.del(2)
|
||||
doAssert(heap.toSortedSeq == @[1, 2])
|
||||
block: # test del last
|
||||
var heap = initHeapQueue[int]()
|
||||
let data = [1, 2, 3]
|
||||
for item in data: push(heap, item)
|
||||
|
||||
heap.del(1)
|
||||
doAssert(heap.toSortedSeq == @[1])
|
||||
heap.del(2)
|
||||
doAssert(heap.toSortedSeq == @[1, 2])
|
||||
|
||||
heap.del(0)
|
||||
doAssert(heap.toSortedSeq == @[])
|
||||
heap.del(1)
|
||||
doAssert(heap.toSortedSeq == @[1])
|
||||
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user