1#
2#
3#            Nim's Runtime Library
4#        (c) Copyright 2016 Yuriy Glukhov
5#
6#    See the file "copying.txt", included in this
7#    distribution, for details about the copyright.
8
9
10## The `heapqueue` module implements a
11## `binary heap data structure<https://en.wikipedia.org/wiki/Binary_heap>`_
12## that can be used as a `priority queue<https://en.wikipedia.org/wiki/Priority_queue>`_.
13## They are represented as arrays for which `a[k] <= a[2*k+1]` and `a[k] <= a[2*k+2]`
14## for all indices `k` (counting elements from 0). The interesting property of a heap is that
15## `a[0]` is always its smallest element.
16##
17## Basic usage
18## -----------
19##
20runnableExamples:
21  var heap = [8, 2].toHeapQueue
22  heap.push(5)
23  # the first element is the lowest element
24  assert heap[0] == 2
25  # remove and return the lowest element
26  assert heap.pop() == 2
27  # the lowest element remaining is 5
28  assert heap[0] == 5
29
30## Usage with custom objects
31## -------------------------
32## To use a `HeapQueue` with a custom object, the `<` operator must be
33## implemented.
34
35runnableExamples:
36  type Job = object
37    priority: int
38
39  proc `<`(a, b: Job): bool = a.priority < b.priority
40
41  var jobs = initHeapQueue[Job]()
42  jobs.push(Job(priority: 1))
43  jobs.push(Job(priority: 2))
44
45  assert jobs[0].priority == 1
46
47
48import std/private/since
49
50type HeapQueue*[T] = object
51  ## A heap queue, commonly known as a priority queue.
52  data: seq[T]
53
54proc initHeapQueue*[T](): HeapQueue[T] =
55  ## Creates a new empty heap.
56  ##
57  ## Heaps are initialized by default, so it is not necessary to call
58  ## this function explicitly.
59  ##
60  ## **See also:**
61  ## * `toHeapQueue proc <#toHeapQueue,openArray[T]>`_
62  discard
63
64proc len*[T](heap: HeapQueue[T]): int {.inline.} =
65  ## Returns the number of elements of `heap`.
66  runnableExamples:
67    let heap = [9, 5, 8].toHeapQueue
68    assert heap.len == 3
69
70  heap.data.len
71
72proc `[]`*[T](heap: HeapQueue[T], i: Natural): lent T {.inline.} =
73  ## Accesses the i-th element of `heap`.
74  heap.data[i]
75
76proc heapCmp[T](x, y: T): bool {.inline.} = x < y
77
78proc siftup[T](heap: var HeapQueue[T], startpos, p: int) =
79  ## `heap` is a heap at all indices >= `startpos`, except possibly for `p`. `p`
80  ## is the index of a leaf with a possibly out-of-order value. Restores the
81  ## heap invariant.
82  var pos = p
83  let newitem = heap[pos]
84  # Follow the path to the root, moving parents down until finding a place
85  # newitem fits.
86  while pos > startpos:
87    let parentpos = (pos - 1) shr 1
88    let parent = heap[parentpos]
89    if heapCmp(newitem, parent):
90      heap.data[pos] = parent
91      pos = parentpos
92    else:
93      break
94  heap.data[pos] = newitem
95
96proc siftdownToBottom[T](heap: var HeapQueue[T], p: int) =
97  # This is faster when the element should be close to the bottom.
98  let endpos = len(heap)
99  var pos = p
100  let startpos = pos
101  let newitem = heap[pos]
102  # Bubble up the smaller child until hitting a leaf.
103  var childpos = 2 * pos + 1 # leftmost child position
104  while childpos < endpos:
105    # Set childpos to index of smaller child.
106    let rightpos = childpos + 1
107    if rightpos < endpos and not heapCmp(heap[childpos], heap[rightpos]):
108      childpos = rightpos
109    # Move the smaller child up.
110    heap.data[pos] = heap[childpos]
111    pos = childpos
112    childpos = 2 * pos + 1
113  # The leaf at pos is empty now. Put newitem there, and bubble it up
114  # to its final resting place (by sifting its parents down).
115  heap.data[pos] = newitem
116  siftup(heap, startpos, pos)
117
118proc siftdown[T](heap: var HeapQueue[T], p: int) =
119  let endpos = len(heap)
120  var pos = p
121  let newitem = heap[pos]
122  var childpos = 2 * pos + 1
123  while childpos < endpos:
124    let rightpos = childpos + 1
125    if rightpos < endpos and not heapCmp(heap[childpos], heap[rightpos]):
126      childpos = rightpos
127    if not heapCmp(heap[childpos], newitem):
128      break
129    heap.data[pos] = heap[childpos]
130    pos = childpos
131    childpos = 2 * pos + 1
132  heap.data[pos] = newitem
133
134proc push*[T](heap: var HeapQueue[T], item: sink T) =
135  ## Pushes `item` onto `heap`, maintaining the heap invariant.
136  heap.data.add(item)
137  siftup(heap, 0, len(heap) - 1)
138
139proc toHeapQueue*[T](x: openArray[T]): HeapQueue[T] {.since: (1, 3).} =
140  ## Creates a new HeapQueue that contains the elements of `x`.
141  ##
142  ## **See also:**
143  ## * `initHeapQueue proc <#initHeapQueue>`_
144  runnableExamples:
145    var heap = [9, 5, 8].toHeapQueue
146    assert heap.pop() == 5
147    assert heap[0] == 8
148
149  # see https://en.wikipedia.org/wiki/Binary_heap#Building_a_heap
150  result.data = @x
151  for i in countdown(x.len div 2 - 1, 0):
152    siftdown(result, i)
153
154proc pop*[T](heap: var HeapQueue[T]): T =
155  ## Pops and returns the smallest item from `heap`,
156  ## maintaining the heap invariant.
157  runnableExamples:
158    var heap = [9, 5, 8].toHeapQueue
159    assert heap.pop() == 5
160
161  let lastelt = heap.data.pop()
162  if heap.len > 0:
163    result = heap[0]
164    heap.data[0] = lastelt
165    siftdownToBottom(heap, 0)
166  else:
167    result = lastelt
168
169proc find*[T](heap: HeapQueue[T], x: T): int {.since: (1, 3).} =
170  ## Linear scan to find the index of the item `x` or -1 if not found.
171  runnableExamples:
172    let heap = [9, 5, 8].toHeapQueue
173    assert heap.find(5) == 0
174    assert heap.find(9) == 1
175    assert heap.find(777) == -1
176
177  result = -1
178  for i in 0 ..< heap.len:
179    if heap[i] == x: return i
180
181proc del*[T](heap: var HeapQueue[T], index: Natural) =
182  ## Removes the element at `index` from `heap`, maintaining the heap invariant.
183  runnableExamples:
184    var heap = [9, 5, 8].toHeapQueue
185    heap.del(1)
186    assert heap[0] == 5
187    assert heap[1] == 8
188
189  swap(heap.data[^1], heap.data[index])
190  let newLen = heap.len - 1
191  heap.data.setLen(newLen)
192  if index < newLen:
193    siftdownToBottom(heap, index)
194
195proc replace*[T](heap: var HeapQueue[T], item: sink T): T =
196  ## Pops and returns the current smallest value, and add the new item.
197  ## This is more efficient than `pop()` followed by `push()`, and can be
198  ## more appropriate when using a fixed-size heap. Note that the value
199  ## returned may be larger than `item`! That constrains reasonable uses of
200  ## this routine unless written as part of a conditional replacement.
201  ##
202  ## **See also:**
203  ## * `pushpop proc <#pushpop,HeapQueue[T],sinkT>`_
204  runnableExamples:
205    var heap = [5, 12].toHeapQueue
206    assert heap.replace(6) == 5
207    assert heap.len == 2
208    assert heap[0] == 6
209    assert heap.replace(4) == 6
210
211  result = heap[0]
212  heap.data[0] = item
213  siftdown(heap, 0)
214
215proc pushpop*[T](heap: var HeapQueue[T], item: sink T): T =
216  ## Fast version of a `push()` followed by a `pop()`.
217  ##
218  ## **See also:**
219  ## * `replace proc <#replace,HeapQueue[T],sinkT>`_
220  runnableExamples:
221    var heap = [5, 12].toHeapQueue
222    assert heap.pushpop(6) == 5
223    assert heap.len == 2
224    assert heap[0] == 6
225    assert heap.pushpop(4) == 4
226
227  result = item
228  if heap.len > 0 and heapCmp(heap.data[0], result):
229    swap(result, heap.data[0])
230    siftdown(heap, 0)
231
232proc clear*[T](heap: var HeapQueue[T]) =
233  ## Removes all elements from `heap`, making it empty.
234  runnableExamples:
235    var heap = [9, 5, 8].toHeapQueue
236    heap.clear()
237    assert heap.len == 0
238
239  heap.data.setLen(0)
240
241proc `$`*[T](heap: HeapQueue[T]): string =
242  ## Turns a heap into its string representation.
243  runnableExamples:
244    let heap = [1, 2].toHeapQueue
245    assert $heap == "[1, 2]"
246
247  result = "["
248  for x in heap.data:
249    if result.len > 1: result.add(", ")
250    result.addQuoted(x)
251  result.add("]")
252