diff options
Diffstat (limited to 'lib/pure/collections/heapqueue.nim')
| -rw-r--r-- | lib/pure/collections/heapqueue.nim | 157 |
1 files changed, 115 insertions, 42 deletions
diff --git a/lib/pure/collections/heapqueue.nim b/lib/pure/collections/heapqueue.nim index 60869142e..cdb8db6e1 100644 --- a/lib/pure/collections/heapqueue.nim +++ b/lib/pure/collections/heapqueue.nim @@ -1,4 +1,3 @@ - # # # Nim's Runtime Library @@ -7,32 +6,74 @@ # See the file "copying.txt", included in this # distribution, for details about the copyright. -##[ Heap queue algorithm (a.k.a. priority queue). Ported from Python heapq. - -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. For the sake of comparison, -non-existing elements are considered to be infinite. The interesting -property of a heap is that a[0] is always its smallest element. - +##[ + 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 + `a[0]` is always its smallest element. + + Basic usage + ----------- + .. code-block:: Nim + import heapqueue + + var heap = initHeapQueue[int]() + heap.push(8) + heap.push(2) + heap.push(5) + # The first element is the lowest element + assert heap[0] == 2 + # Remove and return the lowest element + assert heap.pop() == 2 + # The lowest element remaining is 5 + assert heap[0] == 5 + + Usage with custom object + ------------------------ + To use a `HeapQueue` with a custom object, the `<` operator must be + implemented. + + .. code-block:: Nim + import heapqueue + + type Job = object + priority: int + + proc `<`(a, b: Job): bool = a.priority < b.priority + + var jobs = initHeapQueue[Job]() + jobs.push(Job(priority: 1)) + jobs.push(Job(priority: 2)) + + assert jobs[0].priority == 1 ]## -type HeapQueue*[T] = distinct seq[T] +type HeapQueue*[T] = object + ## A heap queue, commonly known as a priority queue. + data: seq[T] -proc newHeapQueue*[T](): HeapQueue[T] {.inline.} = HeapQueue[T](newSeq[T]()) -proc newHeapQueue*[T](h: var HeapQueue[T]) {.inline.} = h = HeapQueue[T](newSeq[T]()) +proc initHeapQueue*[T](): HeapQueue[T] = + ## Create a new empty heap. + discard -proc len*[T](h: HeapQueue[T]): int {.inline.} = seq[T](h).len -proc `[]`*[T](h: HeapQueue[T], i: int): T {.inline.} = seq[T](h)[i] -proc `[]=`[T](h: var HeapQueue[T], i: int, v: T) {.inline.} = seq[T](h)[i] = v -proc add[T](h: var HeapQueue[T], v: T) {.inline.} = seq[T](h).add(v) +proc len*[T](heap: HeapQueue[T]): int {.inline.} = + ## Return the number of elements of `heap`. + heap.data.len + +proc `[]`*[T](heap: HeapQueue[T], i: Natural): T {.inline.} = + ## Access the i-th element of `heap`. + heap.data[i] proc heapCmp[T](x, y: T): bool {.inline.} = return (x < y) -# 'heap' is a heap at all indices >= startpos, except possibly for pos. pos -# is the index of a leaf with a possibly out-of-order value. Restore the -# heap invariant. proc siftdown[T](heap: var HeapQueue[T], startpos, p: int) = + ## 'heap' is a heap at all indices >= startpos, except possibly for pos. pos + ## is the index of a leaf with a possibly out-of-order value. Restore the + ## heap invariant. var pos = p var newitem = heap[pos] # Follow the path to the root, moving parents down until finding a place @@ -41,11 +82,11 @@ proc siftdown[T](heap: var HeapQueue[T], startpos, p: int) = let parentpos = (pos - 1) shr 1 let parent = heap[parentpos] if heapCmp(newitem, parent): - heap[pos] = parent + heap.data[pos] = parent pos = parentpos else: break - heap[pos] = newitem + heap.data[pos] = newitem proc siftup[T](heap: var HeapQueue[T], p: int) = let endpos = len(heap) @@ -60,48 +101,50 @@ proc siftup[T](heap: var HeapQueue[T], p: int) = if rightpos < endpos and not heapCmp(heap[childpos], heap[rightpos]): childpos = rightpos # Move the smaller child up. - heap[pos] = heap[childpos] + 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 # to its final resting place (by sifting its parents down). - heap[pos] = newitem + heap.data[pos] = newitem siftdown(heap, startpos, pos) proc push*[T](heap: var HeapQueue[T], item: T) = - ## Push item onto heap, maintaining the heap invariant. - (seq[T](heap)).add(item) + ## Push `item` onto heap, maintaining the heap invariant. + heap.data.add(item) siftdown(heap, 0, len(heap)-1) proc pop*[T](heap: var HeapQueue[T]): T = - ## Pop the smallest item off the heap, maintaining the heap invariant. - let lastelt = seq[T](heap).pop() + ## Pop and return the smallest item from `heap`, + ## maintaining the heap invariant. + let lastelt = heap.data.pop() if heap.len > 0: result = heap[0] - heap[0] = lastelt + heap.data[0] = lastelt siftup(heap, 0) else: result = lastelt -proc del*[T](heap: var HeapQueue[T], index: int) = - ## Removes element at `index`, maintaining the heap invariant. - swap(seq[T](heap)[^1], seq[T](heap)[index]) +proc del*[T](heap: var HeapQueue[T], index: Natural) = + ## Removes the element at `index` from `heap`, maintaining the heap invariant. + swap(heap.data[^1], heap.data[index]) let newLen = heap.len - 1 - seq[T](heap).setLen(newLen) + heap.data.setLen(newLen) if index < newLen: heap.siftup(index) proc replace*[T](heap: var HeapQueue[T], item: T): T = ## Pop and return the current smallest value, and add the new item. ## 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 + ## 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: - + ## + ## .. code-block:: nim ## if item > heap[0]: ## item = replace(heap, item) result = heap[0] - heap[0] = item + heap.data[0] = item siftup(heap, 0) proc pushpop*[T](heap: var HeapQueue[T], item: T): T = @@ -111,6 +154,36 @@ proc pushpop*[T](heap: var HeapQueue[T], item: T): T = siftup(heap, 0) return item +proc clear*[T](heap: var HeapQueue[T]) = + ## Remove all elements from `heap`, making it empty. + runnableExamples: + var heap = initHeapQueue[int]() + heap.push(1) + heap.clear() + assert heap.len == 0 + heap.data.setLen(0) + +proc `$`*[T](heap: HeapQueue[T]): string = + ## Turn a heap into its string representation. + runnableExamples: + var heap = initHeapQueue[int]() + heap.push(1) + heap.push(2) + assert $heap == "[1, 2]" + result = "[" + for x in heap.data: + if result.len > 1: result.add(", ") + result.addQuoted(x) + result.add("]") + +proc newHeapQueue*[T](): HeapQueue[T] {.deprecated.} = + ## **Deprecated since v0.20.0:** use ``initHeapQueue`` instead. + initHeapQueue[T]() + +proc newHeapQueue*[T](heap: var HeapQueue[T]) {.deprecated.} = + ## **Deprecated since v0.20.0:** use ``clear`` instead. + heap.clear() + when isMainModule: proc toSortedSeq[T](h: HeapQueue[T]): seq[T] = var tmp = h @@ -119,7 +192,7 @@ when isMainModule: result.add(pop(tmp)) block: # Simple sanity test - var heap = newHeapQueue[int]() + var heap = initHeapQueue[int]() let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0] for item in data: push(heap, item) @@ -127,27 +200,27 @@ when isMainModule: doAssert(heap.toSortedSeq == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) block: # Test del - var heap = newHeapQueue[int]() + 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(0) doAssert(heap[0] == 1) - heap.del(seq[int](heap).find(7)) + heap.del(heap.data.find(7)) doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 5, 6, 8, 9]) - heap.del(seq[int](heap).find(5)) + heap.del(heap.data.find(5)) doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 6, 8, 9]) - heap.del(seq[int](heap).find(6)) + heap.del(heap.data.find(6)) doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 8, 9]) - heap.del(seq[int](heap).find(2)) + heap.del(heap.data.find(2)) doAssert(heap.toSortedSeq == @[1, 3, 4, 8, 9]) block: # Test del last - var heap = newHeapQueue[int]() + var heap = initHeapQueue[int]() let data = [1, 2, 3] for item in data: push(heap, item) |