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