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