diff options
Diffstat (limited to 'lib/pure/collections/heapqueue.nim')
| -rw-r--r-- | lib/pure/collections/heapqueue.nim | 320 |
1 files changed, 176 insertions, 144 deletions
diff --git a/lib/pure/collections/heapqueue.nim b/lib/pure/collections/heapqueue.nim index f958a5b0a..96f9b4430 100644 --- a/lib/pure/collections/heapqueue.nim +++ b/lib/pure/collections/heapqueue.nim @@ -6,76 +6,91 @@ # See the file "copying.txt", included in this # distribution, for details about the copyright. -##[ - 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 -]## +## The `heapqueue` module implements a +## `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 = [8, 2].toHeapQueue + 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 objects +## ------------------------- +## To use a `HeapQueue` with a custom object, the `<` operator must be +## implemented. + +runnableExamples: + 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 + + +import std/private/since + +when defined(nimPreviewSlimSystem): + import std/assertions type HeapQueue*[T] = object ## A heap queue, commonly known as a priority queue. data: seq[T] proc initHeapQueue*[T](): HeapQueue[T] = - ## Create a new empty heap. - discard + ## Creates a new empty heap. + ## + ## Heaps are initialized by default, so it is not necessary to call + ## this function explicitly. + ## + ## **See also:** + ## * `toHeapQueue proc <#toHeapQueue,openArray[T]>`_ + result = default(HeapQueue[T]) proc len*[T](heap: HeapQueue[T]): int {.inline.} = - ## Return the number of elements of `heap`. + ## 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): T {.inline.} = - ## Access the i-th element of `heap`. +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) +iterator items*[T](heap: HeapQueue[T]): lent T {.inline, since: (2, 1, 1).} = + ## Iterates over each item of `heap`. + let L = len(heap) + for i in 0 .. high(heap.data): + yield heap.data[i] + assert(len(heap) == L, "the length of the HeapQueue changed while iterating over it") -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 +proc heapCmp[T](x, y: T): bool {.inline.} = x < y + +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: @@ -88,13 +103,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 @@ -103,132 +119,148 @@ 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: T) = - ## Push `item` onto heap, maintaining the heap invariant. +proc push*[T](heap: var HeapQueue[T], item: sink T) = + ## 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:** + ## * `initHeapQueue proc <#initHeapQueue>`_ + runnableExamples: + var heap = [9, 5, 8].toHeapQueue + assert heap.pop() == 5 + assert heap[0] == 8 + + # 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 = - ## Pop and return the smallest item from `heap`, + ## Pops and returns the smallest item from `heap`, ## maintaining the heap invariant. + runnableExamples: + 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 the index of the item `x` or -1 if not found. + runnableExamples: + 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 + +proc contains*[T](heap: HeapQueue[T], x: T): bool {.since: (2, 1, 1).} = + ## Returns true if `x` is in `heap` or false if not found. This is a shortcut + ## for `find(heap, x) >= 0`. + result = find(heap, x) >= 0 + proc del*[T](heap: var HeapQueue[T], index: Natural) = ## Removes the element at `index` from `heap`, maintaining the heap invariant. + runnableExamples: + 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: 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 +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 ## 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. ## - ## .. code-block:: nim - ## if item > heap[0]: - ## item = replace(heap, item) + ## **See also:** + ## * `pushpop proc <#pushpop,HeapQueue[T],sinkT>`_ + runnableExamples: + 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: T): T = - ## Fast version of a push followed by a pop. - if heap.len > 0 and heapCmp(heap[0], item): - swap(item, heap[0]) - siftup(heap, 0) - return item +proc pushpop*[T](heap: var HeapQueue[T], item: sink T): T = + ## Fast version of a `push()` followed by a `pop()`. + ## + ## **See also:** + ## * `replace proc <#replace,HeapQueue[T],sinkT>`_ + runnableExamples: + 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]) + siftdown(heap, 0) proc clear*[T](heap: var HeapQueue[T]) = - ## Remove all elements from `heap`, making it empty. + ## 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 = - ## Turn a heap into its string representation. + ## 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(", ") 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 - result = @[] - 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[0] == 0) - doAssert(heap.toSortedSeq == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) - - 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(0) - doAssert(heap[0] == 1) - - heap.del(heap.data.find(7)) - doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 5, 6, 8, 9]) - - heap.del(heap.data.find(5)) - doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 6, 8, 9]) - - heap.del(heap.data.find(6)) - doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 8, 9]) - - heap.del(heap.data.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) - - heap.del(2) - doAssert(heap.toSortedSeq == @[1, 2]) - - heap.del(1) - doAssert(heap.toSortedSeq == @[1]) - - heap.del(0) - doAssert(heap.toSortedSeq == @[]) |