Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 1 | :mod:`heapq` --- Heap queue algorithm |
| 2 | ===================================== |
| 3 | |
| 4 | .. module:: heapq |
| 5 | :synopsis: Heap queue algorithm (a.k.a. priority queue). |
| 6 | .. moduleauthor:: Kevin O'Connor |
| 7 | .. sectionauthor:: Guido van Rossum <guido@python.org> |
| 8 | .. sectionauthor:: François Pinard |
Raymond Hettinger | 0a9b518 | 2010-08-07 23:33:05 +0000 | [diff] [blame] | 9 | .. sectionauthor:: Raymond Hettinger |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 10 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 11 | This module provides an implementation of the heap queue algorithm, also known |
| 12 | as the priority queue algorithm. |
| 13 | |
| 14 | Heaps are arrays for which ``heap[k] <= heap[2*k+1]`` and ``heap[k] <= |
| 15 | heap[2*k+2]`` for all *k*, counting elements from zero. For the sake of |
| 16 | comparison, non-existing elements are considered to be infinite. The |
| 17 | interesting property of a heap is that ``heap[0]`` is always its smallest |
| 18 | element. |
| 19 | |
| 20 | The API below differs from textbook heap algorithms in two aspects: (a) We use |
| 21 | zero-based indexing. This makes the relationship between the index for a node |
| 22 | and the indexes for its children slightly less obvious, but is more suitable |
| 23 | since Python uses zero-based indexing. (b) Our pop method returns the smallest |
| 24 | item, not the largest (called a "min heap" in textbooks; a "max heap" is more |
| 25 | common in texts because of its suitability for in-place sorting). |
| 26 | |
| 27 | These two make it possible to view the heap as a regular Python list without |
| 28 | surprises: ``heap[0]`` is the smallest item, and ``heap.sort()`` maintains the |
| 29 | heap invariant! |
| 30 | |
| 31 | To create a heap, use a list initialized to ``[]``, or you can transform a |
| 32 | populated list into a heap via function :func:`heapify`. |
| 33 | |
| 34 | The following functions are provided: |
| 35 | |
| 36 | |
| 37 | .. function:: heappush(heap, item) |
| 38 | |
| 39 | Push the value *item* onto the *heap*, maintaining the heap invariant. |
| 40 | |
| 41 | |
| 42 | .. function:: heappop(heap) |
| 43 | |
| 44 | Pop and return the smallest item from the *heap*, maintaining the heap |
| 45 | invariant. If the heap is empty, :exc:`IndexError` is raised. |
| 46 | |
Benjamin Peterson | 35e8c46 | 2008-04-24 02:34:53 +0000 | [diff] [blame] | 47 | |
Christian Heimes | dd15f6c | 2008-03-16 00:07:10 +0000 | [diff] [blame] | 48 | .. function:: heappushpop(heap, item) |
| 49 | |
| 50 | Push *item* on the heap, then pop and return the smallest item from the |
| 51 | *heap*. The combined action runs more efficiently than :func:`heappush` |
| 52 | followed by a separate call to :func:`heappop`. |
| 53 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 54 | |
| 55 | .. function:: heapify(x) |
| 56 | |
| 57 | Transform list *x* into a heap, in-place, in linear time. |
| 58 | |
| 59 | |
| 60 | .. function:: heapreplace(heap, item) |
| 61 | |
| 62 | Pop and return the smallest item from the *heap*, and also push the new *item*. |
| 63 | The heap size doesn't change. If the heap is empty, :exc:`IndexError` is raised. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 64 | |
Raymond Hettinger | fa16e2c | 2010-09-01 21:26:16 +0000 | [diff] [blame] | 65 | This one step operation is more efficient than a :func:`heappop` followed by |
| 66 | :func:`heappush` and can be more appropriate when using a fixed-size heap. |
| 67 | The pop/push combination always returns an element from the heap and replaces |
| 68 | it with *item*. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 69 | |
Raymond Hettinger | fa16e2c | 2010-09-01 21:26:16 +0000 | [diff] [blame] | 70 | The value returned may be larger than the *item* added. If that isn't |
| 71 | desired, consider using :func:`heappushpop` instead. Its push/pop |
| 72 | combination returns the smaller of the two values, leaving the larger value |
| 73 | on the heap. |
Georg Brandl | af265f4 | 2008-12-07 15:06:20 +0000 | [diff] [blame] | 74 | |
Georg Brandl | 48310cd | 2009-01-03 21:18:54 +0000 | [diff] [blame] | 75 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 76 | The module also offers three general purpose functions based on heaps. |
| 77 | |
| 78 | |
| 79 | .. function:: merge(*iterables) |
| 80 | |
| 81 | Merge multiple sorted inputs into a single sorted output (for example, merge |
Georg Brandl | 9afde1c | 2007-11-01 20:32:30 +0000 | [diff] [blame] | 82 | timestamped entries from multiple log files). Returns an :term:`iterator` |
Benjamin Peterson | 206e307 | 2008-10-19 14:07:49 +0000 | [diff] [blame] | 83 | over the sorted values. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 84 | |
| 85 | Similar to ``sorted(itertools.chain(*iterables))`` but returns an iterable, does |
| 86 | not pull the data into memory all at once, and assumes that each of the input |
| 87 | streams is already sorted (smallest to largest). |
| 88 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 89 | |
Georg Brandl | 036490d | 2009-05-17 13:00:36 +0000 | [diff] [blame] | 90 | .. function:: nlargest(n, iterable, key=None) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 91 | |
| 92 | Return a list with the *n* largest elements from the dataset defined by |
| 93 | *iterable*. *key*, if provided, specifies a function of one argument that is |
| 94 | used to extract a comparison key from each element in the iterable: |
| 95 | ``key=str.lower`` Equivalent to: ``sorted(iterable, key=key, |
| 96 | reverse=True)[:n]`` |
| 97 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 98 | |
Georg Brandl | 036490d | 2009-05-17 13:00:36 +0000 | [diff] [blame] | 99 | .. function:: nsmallest(n, iterable, key=None) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 100 | |
| 101 | Return a list with the *n* smallest elements from the dataset defined by |
| 102 | *iterable*. *key*, if provided, specifies a function of one argument that is |
| 103 | used to extract a comparison key from each element in the iterable: |
| 104 | ``key=str.lower`` Equivalent to: ``sorted(iterable, key=key)[:n]`` |
| 105 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 106 | |
| 107 | The latter two functions perform best for smaller values of *n*. For larger |
| 108 | values, it is more efficient to use the :func:`sorted` function. Also, when |
Georg Brandl | c5605df | 2009-08-13 08:26:44 +0000 | [diff] [blame] | 109 | ``n==1``, it is more efficient to use the built-in :func:`min` and :func:`max` |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 110 | functions. |
| 111 | |
| 112 | |
Raymond Hettinger | fa16e2c | 2010-09-01 21:26:16 +0000 | [diff] [blame] | 113 | Basic Examples |
| 114 | -------------- |
| 115 | |
| 116 | A `heapsort <http://en.wikipedia.org/wiki/Heapsort>`_ can be implemented by |
| 117 | pushing all values onto a heap and then popping off the smallest values one at a |
| 118 | time:: |
| 119 | |
| 120 | >>> def heapsort(iterable): |
| 121 | ... 'Equivalent to sorted(iterable)' |
| 122 | ... h = [] |
| 123 | ... for value in iterable: |
| 124 | ... heappush(h, value) |
| 125 | ... return [heappop(h) for i in range(len(h))] |
| 126 | ... |
| 127 | >>> heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0]) |
| 128 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] |
| 129 | |
| 130 | Heap elements can be tuples. This is useful for assigning comparison values |
| 131 | (such as task priorities) alongside the main record being tracked:: |
| 132 | |
| 133 | >>> h = [] |
| 134 | >>> heappush(h, (5, 'write code')) |
| 135 | >>> heappush(h, (7, 'release product')) |
| 136 | >>> heappush(h, (1, 'write spec')) |
| 137 | >>> heappush(h, (3, 'create tests')) |
| 138 | >>> heappop(h) |
| 139 | (1, 'write spec') |
| 140 | |
| 141 | |
Raymond Hettinger | 0a9b518 | 2010-08-07 23:33:05 +0000 | [diff] [blame] | 142 | Priority Queue Implementation Notes |
| 143 | ----------------------------------- |
| 144 | |
| 145 | A `priority queue <http://en.wikipedia.org/wiki/Priority_queue>`_ is common use |
| 146 | for a heap, and it presents several implementation challenges: |
| 147 | |
| 148 | * Sort stability: how do you get two tasks with equal priorities to be returned |
| 149 | in the order they were originally added? |
| 150 | |
| 151 | * Tuple comparison breaks for (priority, task) pairs if the priorities are equal |
| 152 | and the tasks do not have a default comparison order. |
| 153 | |
Raymond Hettinger | a2156c0 | 2010-08-07 23:38:16 +0000 | [diff] [blame] | 154 | * If the priority of a task changes, how do you move it to a new position in |
Raymond Hettinger | 0a9b518 | 2010-08-07 23:33:05 +0000 | [diff] [blame] | 155 | the heap? |
| 156 | |
| 157 | * Or if a pending task needs to be deleted, how do you find it and remove it |
| 158 | from the queue? |
| 159 | |
| 160 | A solution to the first two challenges is to store entries as 3-element list |
| 161 | including the priority, an entry count, and the task. The entry count serves as |
| 162 | a tie-breaker so that two tasks with the same priority are returned in the order |
| 163 | they were added. And since no two entry counts are the same, the tuple |
| 164 | comparison will never attempt to directly compare two tasks. |
| 165 | |
| 166 | The remaining challenges revolve around finding a pending task and making |
| 167 | changes to its priority or removing it entirely. Finding a task can be done |
| 168 | with a dictionary pointing to an entry in the queue. |
| 169 | |
| 170 | Removing the entry or changing its priority is more difficult because it would |
| 171 | break the heap structure invariants. So, a possible solution is to mark an |
| 172 | entry as invalid and optionally add a new entry with the revised priority:: |
| 173 | |
| 174 | pq = [] # the priority queue list |
| 175 | counter = itertools.count(1) # unique sequence count |
| 176 | task_finder = {} # mapping of tasks to entries |
| 177 | INVALID = 0 # mark an entry as deleted |
| 178 | |
| 179 | def add_task(priority, task, count=None): |
| 180 | if count is None: |
| 181 | count = next(counter) |
| 182 | entry = [priority, count, task] |
| 183 | task_finder[task] = entry |
| 184 | heappush(pq, entry) |
| 185 | |
| 186 | def get_top_priority(): |
| 187 | while True: |
| 188 | priority, count, task = heappop(pq) |
| 189 | del task_finder[task] |
| 190 | if count is not INVALID: |
| 191 | return task |
| 192 | |
| 193 | def delete_task(task): |
| 194 | entry = task_finder[task] |
| 195 | entry[1] = INVALID |
| 196 | |
| 197 | def reprioritize(priority, task): |
| 198 | entry = task_finder[task] |
| 199 | add_task(priority, task, entry[1]) |
| 200 | entry[1] = INVALID |
| 201 | |
| 202 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 203 | Theory |
| 204 | ------ |
| 205 | |
| 206 | (This explanation is due to François Pinard. The Python code for this module |
| 207 | was contributed by Kevin O'Connor.) |
| 208 | |
| 209 | Heaps are arrays for which ``a[k] <= a[2*k+1]`` and ``a[k] <= a[2*k+2]`` for all |
| 210 | *k*, counting elements from 0. For the sake of comparison, non-existing |
| 211 | elements are considered to be infinite. The interesting property of a heap is |
| 212 | that ``a[0]`` is always its smallest element. |
| 213 | |
| 214 | The strange invariant above is meant to be an efficient memory representation |
| 215 | for a tournament. The numbers below are *k*, not ``a[k]``:: |
| 216 | |
| 217 | 0 |
| 218 | |
| 219 | 1 2 |
| 220 | |
| 221 | 3 4 5 6 |
| 222 | |
| 223 | 7 8 9 10 11 12 13 14 |
| 224 | |
| 225 | 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
| 226 | |
| 227 | In the tree above, each cell *k* is topping ``2*k+1`` and ``2*k+2``. In an usual |
| 228 | binary tournament we see in sports, each cell is the winner over the two cells |
| 229 | it tops, and we can trace the winner down the tree to see all opponents s/he |
| 230 | had. However, in many computer applications of such tournaments, we do not need |
| 231 | to trace the history of a winner. To be more memory efficient, when a winner is |
| 232 | promoted, we try to replace it by something else at a lower level, and the rule |
| 233 | becomes that a cell and the two cells it tops contain three different items, but |
| 234 | the top cell "wins" over the two topped cells. |
| 235 | |
| 236 | If this heap invariant is protected at all time, index 0 is clearly the overall |
| 237 | winner. The simplest algorithmic way to remove it and find the "next" winner is |
| 238 | to move some loser (let's say cell 30 in the diagram above) into the 0 position, |
| 239 | and then percolate this new 0 down the tree, exchanging values, until the |
| 240 | invariant is re-established. This is clearly logarithmic on the total number of |
| 241 | items in the tree. By iterating over all items, you get an O(n log n) sort. |
| 242 | |
| 243 | A nice feature of this sort is that you can efficiently insert new items while |
| 244 | the sort is going on, provided that the inserted items are not "better" than the |
| 245 | last 0'th element you extracted. This is especially useful in simulation |
| 246 | contexts, where the tree holds all incoming events, and the "win" condition |
| 247 | means the smallest scheduled time. When an event schedule other events for |
| 248 | execution, they are scheduled into the future, so they can easily go into the |
| 249 | heap. So, a heap is a good structure for implementing schedulers (this is what |
| 250 | I used for my MIDI sequencer :-). |
| 251 | |
| 252 | Various structures for implementing schedulers have been extensively studied, |
| 253 | and heaps are good for this, as they are reasonably speedy, the speed is almost |
| 254 | constant, and the worst case is not much different than the average case. |
| 255 | However, there are other representations which are more efficient overall, yet |
| 256 | the worst cases might be terrible. |
| 257 | |
| 258 | Heaps are also very useful in big disk sorts. You most probably all know that a |
| 259 | big sort implies producing "runs" (which are pre-sorted sequences, which size is |
| 260 | usually related to the amount of CPU memory), followed by a merging passes for |
| 261 | these runs, which merging is often very cleverly organised [#]_. It is very |
| 262 | important that the initial sort produces the longest runs possible. Tournaments |
| 263 | are a good way to that. If, using all the memory available to hold a |
| 264 | tournament, you replace and percolate items that happen to fit the current run, |
| 265 | you'll produce runs which are twice the size of the memory for random input, and |
| 266 | much better for input fuzzily ordered. |
| 267 | |
| 268 | Moreover, if you output the 0'th item on disk and get an input which may not fit |
| 269 | in the current tournament (because the value "wins" over the last output value), |
| 270 | it cannot fit in the heap, so the size of the heap decreases. The freed memory |
| 271 | could be cleverly reused immediately for progressively building a second heap, |
| 272 | which grows at exactly the same rate the first heap is melting. When the first |
| 273 | heap completely vanishes, you switch heaps and start a new run. Clever and |
| 274 | quite effective! |
| 275 | |
| 276 | In a word, heaps are useful memory structures to know. I use them in a few |
| 277 | applications, and I think it is good to keep a 'heap' module around. :-) |
| 278 | |
| 279 | .. rubric:: Footnotes |
| 280 | |
| 281 | .. [#] The disk balancing algorithms which are current, nowadays, are more annoying |
| 282 | than clever, and this is a consequence of the seeking capabilities of the disks. |
| 283 | On devices which cannot seek, like big tape drives, the story was quite |
| 284 | different, and one had to be very clever to ensure (far in advance) that each |
| 285 | tape movement will be the most effective possible (that is, will best |
| 286 | participate at "progressing" the merge). Some tapes were even able to read |
| 287 | backwards, and this was also used to avoid the rewinding time. Believe me, real |
| 288 | good tape sorts were quite spectacular to watch! From all times, sorting has |
| 289 | always been a Great Art! :-) |
| 290 | |