1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
##[ 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.
]##
type HeapQueue*[T] = distinct 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 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 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) =
var pos = p
var newitem = heap[pos]
# Follow the path to the root, moving parents down until finding a place
# newitem fits.
while pos > startpos:
let parentpos = (pos - 1) shr 1
let parent = heap[parentpos]
if heapCmp(newitem, parent):
heap[pos] = parent
pos = parentpos
else:
break
heap[pos] = newitem
proc siftup[T](heap: var HeapQueue[T], p: int) =
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
while childpos < endpos:
# Set childpos to index of smaller child.
let rightpos = childpos + 1
if rightpos < endpos and not heapCmp(heap[childpos], heap[rightpos]):
childpos = rightpos
# Move the smaller child up.
heap[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
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)
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()
if heap.len > 0:
result = heap[0]
heap[0] = lastelt
siftup(heap, 0)
else:
result = lastelt
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
## this routine unless written as part of a conditional replacement:
## if item > heap[0]:
## item = replace(heap, item)
result = heap[0]
heap[0] = item
siftup(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
when isMainModule:
# Simple sanity test
var heap = newHeapQueue[int]()
let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
for item in data:
push(heap, item)
doAssert(heap[0] == 0)
var sort = newSeq[int]()
while heap.len > 0:
sort.add(pop(heap))
doAssert(sort == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
|