mirror of
https://github.com/nim-lang/Nim.git
synced 2026-08-29 18:11:44 +00:00
Cleanup heapqueue module (#10476)
This commit is contained in:
Oscar Nihlgård
committed by
Andreas Rumpf
Andreas Rumpf
parent
2b0ed9995d
commit
9402c82e80
@@ -92,7 +92,7 @@ Collections and algorithms
|
||||
The underlying implementation uses a ``seq``.
|
||||
|
||||
* `heapqueue <heapqueue.html>`_
|
||||
This module implements Heap queue (a.k.a. priority queue) algorithm.
|
||||
Implementation of a heap data structure that can be used as a priority queue.
|
||||
|
||||
* `intsets <intsets.html>`_
|
||||
Efficient implementation of a set of ints as a sparse bit set.
|
||||
|
||||
@@ -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.
|
||||
##[
|
||||
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.
|
||||
|
||||
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.
|
||||
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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user