_heapqmodule.c revision 636488043b490ab70413990b458443aa41f504b2
1/* Drop in replacement for heapq.py 2 3C implementation derived directly from heapq.py in Py2.3 4which was written by Kevin O'Connor, augmented by Tim Peters, 5annotated by François Pinard, and converted to C by Raymond Hettinger. 6 7*/ 8 9#include "Python.h" 10 11static int 12siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos) 13{ 14 PyObject *newitem, *parent; 15 Py_ssize_t parentpos, size; 16 int cmp; 17 18 assert(PyList_Check(heap)); 19 size = PyList_GET_SIZE(heap); 20 if (pos >= size) { 21 PyErr_SetString(PyExc_IndexError, "index out of range"); 22 return -1; 23 } 24 25 /* Follow the path to the root, moving parents down until finding 26 a place newitem fits. */ 27 newitem = PyList_GET_ITEM(heap, pos); 28 while (pos > startpos) { 29 parentpos = (pos - 1) >> 1; 30 parent = PyList_GET_ITEM(heap, parentpos); 31 cmp = PyObject_RichCompareBool(newitem, parent, Py_LT); 32 if (cmp == -1) 33 return -1; 34 if (size != PyList_GET_SIZE(heap)) { 35 PyErr_SetString(PyExc_RuntimeError, 36 "list changed size during iteration"); 37 return -1; 38 } 39 if (cmp == 0) 40 break; 41 parent = PyList_GET_ITEM(heap, parentpos); 42 newitem = PyList_GET_ITEM(heap, pos); 43 PyList_SET_ITEM(heap, parentpos, newitem); 44 PyList_SET_ITEM(heap, pos, parent); 45 pos = parentpos; 46 } 47 return 0; 48} 49 50static int 51siftup(PyListObject *heap, Py_ssize_t pos) 52{ 53 Py_ssize_t startpos, endpos, childpos, rightpos, limit; 54 PyObject *tmp1, *tmp2; 55 int cmp; 56 57 assert(PyList_Check(heap)); 58 endpos = PyList_GET_SIZE(heap); 59 startpos = pos; 60 if (pos >= endpos) { 61 PyErr_SetString(PyExc_IndexError, "index out of range"); 62 return -1; 63 } 64 65 /* Bubble up the smaller child until hitting a leaf. */ 66 limit = endpos / 2; /* smallest pos that has no child */ 67 while (pos < limit) { 68 /* Set childpos to index of smaller child. */ 69 childpos = 2*pos + 1; /* leftmost child position */ 70 rightpos = childpos + 1; 71 if (rightpos < endpos) { 72 cmp = PyObject_RichCompareBool( 73 PyList_GET_ITEM(heap, childpos), 74 PyList_GET_ITEM(heap, rightpos), 75 Py_LT); 76 if (cmp == -1) 77 return -1; 78 if (cmp == 0) 79 childpos = rightpos; 80 if (endpos != PyList_GET_SIZE(heap)) { 81 PyErr_SetString(PyExc_RuntimeError, 82 "list changed size during iteration"); 83 return -1; 84 } 85 } 86 /* Move the smaller child up. */ 87 tmp1 = PyList_GET_ITEM(heap, childpos); 88 tmp2 = PyList_GET_ITEM(heap, pos); 89 PyList_SET_ITEM(heap, childpos, tmp2); 90 PyList_SET_ITEM(heap, pos, tmp1); 91 pos = childpos; 92 } 93 /* Bubble it up to its final resting place (by sifting its parents down). */ 94 return siftdown(heap, startpos, pos); 95} 96 97static PyObject * 98heappush(PyObject *self, PyObject *args) 99{ 100 PyObject *heap, *item; 101 102 if (!PyArg_UnpackTuple(args, "heappush", 2, 2, &heap, &item)) 103 return NULL; 104 105 if (!PyList_Check(heap)) { 106 PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 107 return NULL; 108 } 109 110 if (PyList_Append(heap, item)) 111 return NULL; 112 113 if (siftdown((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-1)) 114 return NULL; 115 Py_RETURN_NONE; 116} 117 118PyDoc_STRVAR(heappush_doc, 119"heappush(heap, item) -> None. Push item onto heap, maintaining the heap invariant."); 120 121static PyObject * 122heappop_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t)) 123{ 124 PyObject *lastelt, *returnitem; 125 Py_ssize_t n; 126 127 if (!PyList_Check(heap)) { 128 PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 129 return NULL; 130 } 131 132 /* raises IndexError if the heap is empty */ 133 n = PyList_GET_SIZE(heap); 134 if (n == 0) { 135 PyErr_SetString(PyExc_IndexError, "index out of range"); 136 return NULL; 137 } 138 139 lastelt = PyList_GET_ITEM(heap, n-1) ; 140 Py_INCREF(lastelt); 141 if (PyList_SetSlice(heap, n-1, n, NULL)) { 142 Py_DECREF(lastelt); 143 return NULL; 144 } 145 n--; 146 147 if (!n) 148 return lastelt; 149 returnitem = PyList_GET_ITEM(heap, 0); 150 PyList_SET_ITEM(heap, 0, lastelt); 151 if (siftup_func((PyListObject *)heap, 0)) { 152 Py_DECREF(returnitem); 153 return NULL; 154 } 155 return returnitem; 156} 157 158static PyObject * 159heappop(PyObject *self, PyObject *heap) 160{ 161 return heappop_internal(heap, siftup); 162} 163 164PyDoc_STRVAR(heappop_doc, 165"Pop the smallest item off the heap, maintaining the heap invariant."); 166 167static PyObject * 168heapreplace_internal(PyObject *args, int siftup_func(PyListObject *, Py_ssize_t)) 169{ 170 PyObject *heap, *item, *returnitem; 171 172 if (!PyArg_UnpackTuple(args, "heapreplace", 2, 2, &heap, &item)) 173 return NULL; 174 175 if (!PyList_Check(heap)) { 176 PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 177 return NULL; 178 } 179 180 if (PyList_GET_SIZE(heap) == 0) { 181 PyErr_SetString(PyExc_IndexError, "index out of range"); 182 return NULL; 183 } 184 185 returnitem = PyList_GET_ITEM(heap, 0); 186 Py_INCREF(item); 187 PyList_SET_ITEM(heap, 0, item); 188 if (siftup_func((PyListObject *)heap, 0)) { 189 Py_DECREF(returnitem); 190 return NULL; 191 } 192 return returnitem; 193} 194 195static PyObject * 196heapreplace(PyObject *self, PyObject *args) 197{ 198 return heapreplace_internal(args, siftup); 199} 200 201PyDoc_STRVAR(heapreplace_doc, 202"heapreplace(heap, item) -> value. Pop and return the current smallest value, and add the new item.\n\ 203\n\ 204This is more efficient than heappop() followed by heappush(), and can be\n\ 205more appropriate when using a fixed-size heap. Note that the value\n\ 206returned may be larger than item! That constrains reasonable uses of\n\ 207this routine unless written as part of a conditional replacement:\n\n\ 208 if item > heap[0]:\n\ 209 item = heapreplace(heap, item)\n"); 210 211static PyObject * 212heappushpop(PyObject *self, PyObject *args) 213{ 214 PyObject *heap, *item, *returnitem; 215 int cmp; 216 217 if (!PyArg_UnpackTuple(args, "heappushpop", 2, 2, &heap, &item)) 218 return NULL; 219 220 if (!PyList_Check(heap)) { 221 PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 222 return NULL; 223 } 224 225 if (PyList_GET_SIZE(heap) == 0) { 226 Py_INCREF(item); 227 return item; 228 } 229 230 cmp = PyObject_RichCompareBool(PyList_GET_ITEM(heap, 0), item, Py_LT); 231 if (cmp == -1) 232 return NULL; 233 if (cmp == 0) { 234 Py_INCREF(item); 235 return item; 236 } 237 238 if (PyList_GET_SIZE(heap) == 0) { 239 PyErr_SetString(PyExc_IndexError, "index out of range"); 240 return NULL; 241 } 242 243 returnitem = PyList_GET_ITEM(heap, 0); 244 Py_INCREF(item); 245 PyList_SET_ITEM(heap, 0, item); 246 if (siftup((PyListObject *)heap, 0)) { 247 Py_DECREF(returnitem); 248 return NULL; 249 } 250 return returnitem; 251} 252 253PyDoc_STRVAR(heappushpop_doc, 254"heappushpop(heap, item) -> value. Push item on the heap, then pop and return the smallest item\n\ 255from the heap. The combined action runs more efficiently than\n\ 256heappush() followed by a separate call to heappop()."); 257 258static Py_ssize_t 259keep_top_bit(Py_ssize_t n) 260{ 261 int i = 0; 262 263 while (n > 1) { 264 i += 1; 265 n >>= 1; 266 } 267 return n << i; 268} 269 270/* Cache friendly version of heapify() 271 ----------------------------------- 272 273 Build-up a heap in O(n) time by performing siftup() operations 274 on nodes whose children are already heaps. 275 276 The simplest way is to sift the nodes in reverse order from 277 n//2-1 to 0 inclusive. The downside is that children may be 278 out of cache by the time their parent is reached. 279 280 A better way is to not wait for the children to go out of cache. 281 Once a sibling pair of child nodes have been sifted, immediately 282 sift their parent node (while the children are still in cache). 283 284 Both ways build child heaps before their parents, so both ways 285 do the exact same number of comparisons and produce exactly 286 the same heap. The only difference is that the traversal 287 order is optimized for cache efficiency. 288*/ 289 290static PyObject * 291cache_friendly_heapify(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t)) 292{ 293 Py_ssize_t i, j, m, mhalf, leftmost; 294 295 m = PyList_GET_SIZE(heap) >> 1; /* index of first childless node */ 296 leftmost = keep_top_bit(m + 1) - 1; /* leftmost node in row of m */ 297 mhalf = m >> 1; /* parent of first childless node */ 298 299 for (i = leftmost - 1 ; i >= mhalf ; i--) { 300 j = i; 301 while (1) { 302 if (siftup_func((PyListObject *)heap, j)) 303 return NULL; 304 if (!(j & 1)) 305 break; 306 j >>= 1; 307 } 308 } 309 310 for (i = m - 1 ; i >= leftmost ; i--) { 311 j = i; 312 while (1) { 313 if (siftup_func((PyListObject *)heap, j)) 314 return NULL; 315 if (!(j & 1)) 316 break; 317 j >>= 1; 318 } 319 } 320 Py_RETURN_NONE; 321} 322 323static PyObject * 324heapify_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t)) 325{ 326 Py_ssize_t i, n; 327 328 if (!PyList_Check(heap)) { 329 PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 330 return NULL; 331 } 332 333 /* For heaps likely to be bigger than L1 cache, we use the cache 334 friendly heapify function. For smaller heaps that fit entirely 335 in cache, we prefer the simpler algorithm with less branching. 336 */ 337 n = PyList_GET_SIZE(heap); 338 if (n > 2500) 339 return cache_friendly_heapify(heap, siftup_func); 340 341 /* Transform bottom-up. The largest index there's any point to 342 looking at is the largest with a child index in-range, so must 343 have 2*i + 1 < n, or i < (n-1)/2. If n is even = 2*j, this is 344 (2*j-1)/2 = j-1/2 so j-1 is the largest, which is n//2 - 1. If 345 n is odd = 2*j+1, this is (2*j+1-1)/2 = j so j-1 is the largest, 346 and that's again n//2-1. 347 */ 348 for (i = n/2 - 1 ; i >= 0 ; i--) 349 if (siftup_func((PyListObject *)heap, i)) 350 return NULL; 351 Py_RETURN_NONE; 352} 353 354static PyObject * 355heapify(PyObject *self, PyObject *heap) 356{ 357 return heapify_internal(heap, siftup); 358} 359 360PyDoc_STRVAR(heapify_doc, 361"Transform list into a heap, in-place, in O(len(heap)) time."); 362 363static int 364siftdown_max(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos) 365{ 366 PyObject *newitem, *parent; 367 Py_ssize_t parentpos, size; 368 int cmp; 369 370 assert(PyList_Check(heap)); 371 size = PyList_GET_SIZE(heap); 372 if (pos >= size) { 373 PyErr_SetString(PyExc_IndexError, "index out of range"); 374 return -1; 375 } 376 377 /* Follow the path to the root, moving parents down until finding 378 a place newitem fits. */ 379 newitem = PyList_GET_ITEM(heap, pos); 380 while (pos > startpos) { 381 parentpos = (pos - 1) >> 1; 382 parent = PyList_GET_ITEM(heap, parentpos); 383 cmp = PyObject_RichCompareBool(parent, newitem, Py_LT); 384 if (cmp == -1) 385 return -1; 386 if (size != PyList_GET_SIZE(heap)) { 387 PyErr_SetString(PyExc_RuntimeError, 388 "list changed size during iteration"); 389 return -1; 390 } 391 if (cmp == 0) 392 break; 393 parent = PyList_GET_ITEM(heap, parentpos); 394 newitem = PyList_GET_ITEM(heap, pos); 395 PyList_SET_ITEM(heap, parentpos, newitem); 396 PyList_SET_ITEM(heap, pos, parent); 397 pos = parentpos; 398 } 399 return 0; 400} 401 402static int 403siftup_max(PyListObject *heap, Py_ssize_t pos) 404{ 405 Py_ssize_t startpos, endpos, childpos, rightpos, limit; 406 PyObject *tmp1, *tmp2; 407 int cmp; 408 409 assert(PyList_Check(heap)); 410 endpos = PyList_GET_SIZE(heap); 411 startpos = pos; 412 if (pos >= endpos) { 413 PyErr_SetString(PyExc_IndexError, "index out of range"); 414 return -1; 415 } 416 417 /* Bubble up the smaller child until hitting a leaf. */ 418 limit = endpos / 2; /* smallest pos that has no child */ 419 while (pos < limit) { 420 /* Set childpos to index of smaller child. */ 421 childpos = 2*pos + 1; /* leftmost child position */ 422 rightpos = childpos + 1; 423 if (rightpos < endpos) { 424 cmp = PyObject_RichCompareBool( 425 PyList_GET_ITEM(heap, rightpos), 426 PyList_GET_ITEM(heap, childpos), 427 Py_LT); 428 if (cmp == -1) 429 return -1; 430 if (cmp == 0) 431 childpos = rightpos; 432 if (endpos != PyList_GET_SIZE(heap)) { 433 PyErr_SetString(PyExc_RuntimeError, 434 "list changed size during iteration"); 435 return -1; 436 } 437 } 438 /* Move the smaller child up. */ 439 tmp1 = PyList_GET_ITEM(heap, childpos); 440 tmp2 = PyList_GET_ITEM(heap, pos); 441 PyList_SET_ITEM(heap, childpos, tmp2); 442 PyList_SET_ITEM(heap, pos, tmp1); 443 pos = childpos; 444 } 445 /* Bubble it up to its final resting place (by sifting its parents down). */ 446 return siftdown_max(heap, startpos, pos); 447} 448 449static PyObject * 450heappop_max(PyObject *self, PyObject *heap) 451{ 452 return heappop_internal(heap, siftup_max); 453} 454 455PyDoc_STRVAR(heappop_max_doc, "Maxheap variant of heappop."); 456 457static PyObject * 458heapreplace_max(PyObject *self, PyObject *args) 459{ 460 return heapreplace_internal(args, siftup_max); 461} 462 463PyDoc_STRVAR(heapreplace_max_doc, "Maxheap variant of heapreplace"); 464 465static PyObject * 466heapify_max(PyObject *self, PyObject *heap) 467{ 468 return heapify_internal(heap, siftup_max); 469} 470 471PyDoc_STRVAR(heapify_max_doc, "Maxheap variant of heapify."); 472 473static PyMethodDef heapq_methods[] = { 474 {"heappush", (PyCFunction)heappush, 475 METH_VARARGS, heappush_doc}, 476 {"heappushpop", (PyCFunction)heappushpop, 477 METH_VARARGS, heappushpop_doc}, 478 {"heappop", (PyCFunction)heappop, 479 METH_O, heappop_doc}, 480 {"heapreplace", (PyCFunction)heapreplace, 481 METH_VARARGS, heapreplace_doc}, 482 {"heapify", (PyCFunction)heapify, 483 METH_O, heapify_doc}, 484 {"_heappop_max", (PyCFunction)heappop_max, 485 METH_O, heappop_max_doc}, 486 {"_heapreplace_max",(PyCFunction)heapreplace_max, 487 METH_VARARGS, heapreplace_max_doc}, 488 {"_heapify_max", (PyCFunction)heapify_max, 489 METH_O, heapify_max_doc}, 490 {NULL, NULL} /* sentinel */ 491}; 492 493PyDoc_STRVAR(module_doc, 494"Heap queue algorithm (a.k.a. priority queue).\n\ 495\n\ 496Heaps are arrays for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for\n\ 497all k, counting elements from 0. For the sake of comparison,\n\ 498non-existing elements are considered to be infinite. The interesting\n\ 499property of a heap is that a[0] is always its smallest element.\n\ 500\n\ 501Usage:\n\ 502\n\ 503heap = [] # creates an empty heap\n\ 504heappush(heap, item) # pushes a new item on the heap\n\ 505item = heappop(heap) # pops the smallest item from the heap\n\ 506item = heap[0] # smallest item on the heap without popping it\n\ 507heapify(x) # transforms list into a heap, in-place, in linear time\n\ 508item = heapreplace(heap, item) # pops and returns smallest item, and adds\n\ 509 # new item; the heap size is unchanged\n\ 510\n\ 511Our API differs from textbook heap algorithms as follows:\n\ 512\n\ 513- We use 0-based indexing. This makes the relationship between the\n\ 514 index for a node and the indexes for its children slightly less\n\ 515 obvious, but is more suitable since Python uses 0-based indexing.\n\ 516\n\ 517- Our heappop() method returns the smallest item, not the largest.\n\ 518\n\ 519These two make it possible to view the heap as a regular Python list\n\ 520without surprises: heap[0] is the smallest item, and heap.sort()\n\ 521maintains the heap invariant!\n"); 522 523 524PyDoc_STRVAR(__about__, 525"Heap queues\n\ 526\n\ 527[explanation by Fran\xc3\xa7ois Pinard]\n\ 528\n\ 529Heaps are arrays for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for\n\ 530all k, counting elements from 0. For the sake of comparison,\n\ 531non-existing elements are considered to be infinite. The interesting\n\ 532property of a heap is that a[0] is always its smallest element.\n" 533"\n\ 534The strange invariant above is meant to be an efficient memory\n\ 535representation for a tournament. The numbers below are `k', not a[k]:\n\ 536\n\ 537 0\n\ 538\n\ 539 1 2\n\ 540\n\ 541 3 4 5 6\n\ 542\n\ 543 7 8 9 10 11 12 13 14\n\ 544\n\ 545 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30\n\ 546\n\ 547\n\ 548In the tree above, each cell `k' is topping `2*k+1' and `2*k+2'. In\n\ 549an usual binary tournament we see in sports, each cell is the winner\n\ 550over the two cells it tops, and we can trace the winner down the tree\n\ 551to see all opponents s/he had. However, in many computer applications\n\ 552of such tournaments, we do not need to trace the history of a winner.\n\ 553To be more memory efficient, when a winner is promoted, we try to\n\ 554replace it by something else at a lower level, and the rule becomes\n\ 555that a cell and the two cells it tops contain three different items,\n\ 556but the top cell \"wins\" over the two topped cells.\n" 557"\n\ 558If this heap invariant is protected at all time, index 0 is clearly\n\ 559the overall winner. The simplest algorithmic way to remove it and\n\ 560find the \"next\" winner is to move some loser (let's say cell 30 in the\n\ 561diagram above) into the 0 position, and then percolate this new 0 down\n\ 562the tree, exchanging values, until the invariant is re-established.\n\ 563This is clearly logarithmic on the total number of items in the tree.\n\ 564By iterating over all items, you get an O(n ln n) sort.\n" 565"\n\ 566A nice feature of this sort is that you can efficiently insert new\n\ 567items while the sort is going on, provided that the inserted items are\n\ 568not \"better\" than the last 0'th element you extracted. This is\n\ 569especially useful in simulation contexts, where the tree holds all\n\ 570incoming events, and the \"win\" condition means the smallest scheduled\n\ 571time. When an event schedule other events for execution, they are\n\ 572scheduled into the future, so they can easily go into the heap. So, a\n\ 573heap is a good structure for implementing schedulers (this is what I\n\ 574used for my MIDI sequencer :-).\n" 575"\n\ 576Various structures for implementing schedulers have been extensively\n\ 577studied, and heaps are good for this, as they are reasonably speedy,\n\ 578the speed is almost constant, and the worst case is not much different\n\ 579than the average case. However, there are other representations which\n\ 580are more efficient overall, yet the worst cases might be terrible.\n" 581"\n\ 582Heaps are also very useful in big disk sorts. You most probably all\n\ 583know that a big sort implies producing \"runs\" (which are pre-sorted\n\ 584sequences, which size is usually related to the amount of CPU memory),\n\ 585followed by a merging passes for these runs, which merging is often\n\ 586very cleverly organised[1]. It is very important that the initial\n\ 587sort produces the longest runs possible. Tournaments are a good way\n\ 588to that. If, using all the memory available to hold a tournament, you\n\ 589replace and percolate items that happen to fit the current run, you'll\n\ 590produce runs which are twice the size of the memory for random input,\n\ 591and much better for input fuzzily ordered.\n" 592"\n\ 593Moreover, if you output the 0'th item on disk and get an input which\n\ 594may not fit in the current tournament (because the value \"wins\" over\n\ 595the last output value), it cannot fit in the heap, so the size of the\n\ 596heap decreases. The freed memory could be cleverly reused immediately\n\ 597for progressively building a second heap, which grows at exactly the\n\ 598same rate the first heap is melting. When the first heap completely\n\ 599vanishes, you switch heaps and start a new run. Clever and quite\n\ 600effective!\n\ 601\n\ 602In a word, heaps are useful memory structures to know. I use them in\n\ 603a few applications, and I think it is good to keep a `heap' module\n\ 604around. :-)\n" 605"\n\ 606--------------------\n\ 607[1] The disk balancing algorithms which are current, nowadays, are\n\ 608more annoying than clever, and this is a consequence of the seeking\n\ 609capabilities of the disks. On devices which cannot seek, like big\n\ 610tape drives, the story was quite different, and one had to be very\n\ 611clever to ensure (far in advance) that each tape movement will be the\n\ 612most effective possible (that is, will best participate at\n\ 613\"progressing\" the merge). Some tapes were even able to read\n\ 614backwards, and this was also used to avoid the rewinding time.\n\ 615Believe me, real good tape sorts were quite spectacular to watch!\n\ 616From all times, sorting has always been a Great Art! :-)\n"); 617 618 619static struct PyModuleDef _heapqmodule = { 620 PyModuleDef_HEAD_INIT, 621 "_heapq", 622 module_doc, 623 -1, 624 heapq_methods, 625 NULL, 626 NULL, 627 NULL, 628 NULL 629}; 630 631PyMODINIT_FUNC 632PyInit__heapq(void) 633{ 634 PyObject *m, *about; 635 636 m = PyModule_Create(&_heapqmodule); 637 if (m == NULL) 638 return NULL; 639 about = PyUnicode_DecodeUTF8(__about__, strlen(__about__), NULL); 640 PyModule_AddObject(m, "__about__", about); 641 return m; 642} 643 644