_heapqmodule.c revision 1a21451b1d73b65af949193208372e86bf308411
144aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera/* Drop in replacement for heapq.py 244aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera 344aed07672d7775f168a49dbb5b8a13682c02600Shubham AjmeraC implementation derived directly from heapq.py in Py2.3 444aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmerawhich was written by Kevin O'Connor, augmented by Tim Peters, 544aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmeraannotated by Fran�ois Pinard, and converted to C by Raymond Hettinger. 644aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera 744aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera*/ 844aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera 944aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera#include "Python.h" 1044aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera 1144aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera/* Older implementations of heapq used Py_LE for comparisons. Now, it uses 1244aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera Py_LT so it will match min(), sorted(), and bisect(). Unfortunately, some 1344aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera client code (Twisted for example) relied on Py_LE, so this little function 1444aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera restores compatability by trying both. 1544aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera*/ 1644aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmerastatic int 1744aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmeracmp_lt(PyObject *x, PyObject *y) 1844aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera{ 199efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath int cmp; 209efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath cmp = PyObject_RichCompareBool(x, y, Py_LT); 219efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (cmp == -1 && PyErr_ExceptionMatches(PyExc_AttributeError)) { 229efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyErr_Clear(); 239efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath cmp = PyObject_RichCompareBool(y, x, Py_LE); 249efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (cmp != -1) 250976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera cmp = 1 - cmp; 2644aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera } 279efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return cmp; 289efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath} 299efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 309efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathstatic int 319efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath_siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos) 3244aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera{ 3344aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera PyObject *newitem, *parent; 3444aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera int cmp; 3544aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera Py_ssize_t parentpos; 3644aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera 379efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath assert(PyList_Check(heap)); 3844aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera if (pos >= PyList_GET_SIZE(heap)) { 3944aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera PyErr_SetString(PyExc_IndexError, "index out of range"); 4044aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera return -1; 4144aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera } 4244aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera 4344aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera newitem = PyList_GET_ITEM(heap, pos); 4444aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera Py_INCREF(newitem); 4544aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera /* Follow the path to the root, moving parents down until finding 4644aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera a place newitem fits. */ 4744aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera while (pos > startpos){ 4844aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera parentpos = (pos - 1) >> 1; 4944aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera parent = PyList_GET_ITEM(heap, parentpos); 5044aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera cmp = cmp_lt(newitem, parent); 5144aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera if (cmp == -1) { 5244aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera Py_DECREF(newitem); 5344aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera return -1; 5444aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera } 5544aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera if (cmp == 0) 5644aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera break; 5744aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera Py_INCREF(parent); 5844aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera Py_DECREF(PyList_GET_ITEM(heap, pos)); 5944aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera PyList_SET_ITEM(heap, pos, parent); 6044aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera pos = parentpos; 6144aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera } 6244aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera Py_DECREF(PyList_GET_ITEM(heap, pos)); 6344aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera PyList_SET_ITEM(heap, pos, newitem); 6444aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera return 0; 6544aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera} 6644aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera 6744aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmerastatic int 6844aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera_siftup(PyListObject *heap, Py_ssize_t pos) 6944aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera{ 7044aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera Py_ssize_t startpos, endpos, childpos, rightpos; 7144aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera int cmp; 7244aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera PyObject *newitem, *tmp; 7344aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera 7444aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera assert(PyList_Check(heap)); 7544aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera endpos = PyList_GET_SIZE(heap); 7644aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera startpos = pos; 7744aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera if (pos >= endpos) { 7844aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera PyErr_SetString(PyExc_IndexError, "index out of range"); 7944aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera return -1; 8044aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera } 8144aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera newitem = PyList_GET_ITEM(heap, pos); 8244aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera Py_INCREF(newitem); 830976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera 849efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath /* Bubble up the smaller child until hitting a leaf. */ 859efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath childpos = 2*pos + 1; /* leftmost child position */ 869efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath while (childpos < endpos) { 879efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath /* Set childpos to index of smaller child. */ 889efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath rightpos = childpos + 1; 899efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (rightpos < endpos) { 909efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath cmp = cmp_lt( 919efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyList_GET_ITEM(heap, childpos), 929efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyList_GET_ITEM(heap, rightpos)); 939efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (cmp == -1) { 949efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_DECREF(newitem); 959efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return -1; 969efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 979efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (cmp == 0) 989efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath childpos = rightpos; 999efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 1009efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath /* Move the smaller child up. */ 1019efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath tmp = PyList_GET_ITEM(heap, childpos); 1029efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_INCREF(tmp); 1039efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_DECREF(PyList_GET_ITEM(heap, pos)); 1049efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyList_SET_ITEM(heap, pos, tmp); 105508f0abe4f8686d29466b2c364cb4f60bf0484b5Tobias Thierer pos = childpos; 1069efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath childpos = 2*pos + 1; 1079efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 1089efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1099efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath /* The leaf at pos is empty now. Put newitem there, and and bubble 1109efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath it up to its final resting place (by sifting its parents down). */ 1119efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_DECREF(PyList_GET_ITEM(heap, pos)); 1129efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyList_SET_ITEM(heap, pos, newitem); 1139efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return _siftdown(heap, startpos, pos); 1149efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath} 1159efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1169efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathstatic PyObject * 1179efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathheappush(PyObject *self, PyObject *args) 1189efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath{ 1199efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyObject *heap, *item; 1209efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1219efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (!PyArg_UnpackTuple(args, "heappush", 2, 2, &heap, &item)) 1229efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1239efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1249efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (!PyList_Check(heap)) { 1259efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 1269efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1279efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 1289efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 129508f0abe4f8686d29466b2c364cb4f60bf0484b5Tobias Thierer if (PyList_Append(heap, item) == -1) 1309efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1319efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1329efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (_siftdown((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-1) == -1) 1339efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1349efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_INCREF(Py_None); 1359efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return Py_None; 1369efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath} 1379efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1389efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan KamathPyDoc_STRVAR(heappush_doc, 1399efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath"Push item onto heap, maintaining the heap invariant."); 1409efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1419efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathstatic PyObject * 1429efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathheappop(PyObject *self, PyObject *heap) 1439efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath{ 1449efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyObject *lastelt, *returnitem; 1459efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_ssize_t n; 1469efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1479efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (!PyList_Check(heap)) { 1489efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 1499efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1509efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 1519efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1529efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath /* # raises appropriate IndexError if heap is empty */ 1539efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath n = PyList_GET_SIZE(heap); 1549efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (n == 0) { 1559efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyErr_SetString(PyExc_IndexError, "index out of range"); 1569efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1579efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 1589efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1599efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath lastelt = PyList_GET_ITEM(heap, n-1) ; 1609efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_INCREF(lastelt); 1619efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyList_SetSlice(heap, n-1, n, NULL); 1629efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath n--; 1639efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1649efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (!n) 1659efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return lastelt; 1669efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath returnitem = PyList_GET_ITEM(heap, 0); 1679efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyList_SET_ITEM(heap, 0, lastelt); 1689efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (_siftup((PyListObject *)heap, 0) == -1) { 1699efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_DECREF(returnitem); 1709efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1719efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 1729efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return returnitem; 1739efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath} 1749efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1759efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan KamathPyDoc_STRVAR(heappop_doc, 1769efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath"Pop the smallest item off the heap, maintaining the heap invariant."); 1779efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1789efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathstatic PyObject * 1799efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathheapreplace(PyObject *self, PyObject *args) 1809efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath{ 1819efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyObject *heap, *item, *returnitem; 1829efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1839efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (!PyArg_UnpackTuple(args, "heapreplace", 2, 2, &heap, &item)) 1849efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1859efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1869efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (!PyList_Check(heap)) { 1879efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 1889efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1899efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 1909efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1919efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (PyList_GET_SIZE(heap) < 1) { 1929efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyErr_SetString(PyExc_IndexError, "index out of range"); 1939efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 1949efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 1959efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 1969efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath returnitem = PyList_GET_ITEM(heap, 0); 1979efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_INCREF(item); 1989efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyList_SET_ITEM(heap, 0, item); 1999efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (_siftup((PyListObject *)heap, 0) == -1) { 2009efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_DECREF(returnitem); 2019efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 2029efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 2039efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return returnitem; 2049efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath} 2059efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 2069efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan KamathPyDoc_STRVAR(heapreplace_doc, 2079efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath"Pop and return the current smallest value, and add the new item.\n\ 2089efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath\n\ 2099efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan KamathThis is more efficient than heappop() followed by heappush(), and can be\n\ 2109efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathmore appropriate when using a fixed-size heap. Note that the value\n\ 2119efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathreturned may be larger than item! That constrains reasonable uses of\n\ 2129efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamaththis routine unless written as part of a conditional replacement:\n\n\ 2139efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if item > heap[0]:\n\ 2149efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath item = heapreplace(heap, item)\n"); 2159efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 2169efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathstatic PyObject * 2179efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamathheappushpop(PyObject *self, PyObject *args) 2189efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath{ 2199efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath PyObject *heap, *item, *returnitem; 2209efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath int cmp; 2219efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 2229efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (!PyArg_UnpackTuple(args, "heappushpop", 2, 2, &heap, &item)) 2239efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return NULL; 2249efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath 2259efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath if (!PyList_Check(heap)) { 2260976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 2270976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera return NULL; 2280976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera } 2290976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera 2300976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera if (PyList_GET_SIZE(heap) < 1) { 2310976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera Py_INCREF(item); 2320976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera return item; 2330976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera } 2340976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera 2350976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera cmp = cmp_lt(PyList_GET_ITEM(heap, 0), item); 2360976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera if (cmp == -1) 2370976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera return NULL; 2380976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera if (cmp == 0) { 2399efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath Py_INCREF(item); 2409efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath return item; 2419efb6d12ce4d2ffedb73d6e9887ea2c89f8ec129Narayan Kamath } 2420976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera 2430976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera returnitem = PyList_GET_ITEM(heap, 0); 2440976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera Py_INCREF(item); 2450976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera PyList_SET_ITEM(heap, 0, item); 2460976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera if (_siftup((PyListObject *)heap, 0) == -1) { 2470976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera Py_DECREF(returnitem); 2480976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera return NULL; 2490976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera } 2500976dc2e109a3ca2bd977d18eee74e4b7c9ada30Shubham Ajmera return returnitem; 25144aed07672d7775f168a49dbb5b8a13682c02600Shubham Ajmera} 252 253PyDoc_STRVAR(heappushpop_doc, 254"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 PyObject * 259heapify(PyObject *self, PyObject *heap) 260{ 261 Py_ssize_t i, n; 262 263 if (!PyList_Check(heap)) { 264 PyErr_SetString(PyExc_TypeError, "heap argument must be a list"); 265 return NULL; 266 } 267 268 n = PyList_GET_SIZE(heap); 269 /* Transform bottom-up. The largest index there's any point to 270 looking at is the largest with a child index in-range, so must 271 have 2*i + 1 < n, or i < (n-1)/2. If n is even = 2*j, this is 272 (2*j-1)/2 = j-1/2 so j-1 is the largest, which is n//2 - 1. If 273 n is odd = 2*j+1, this is (2*j+1-1)/2 = j so j-1 is the largest, 274 and that's again n//2-1. 275 */ 276 for (i=n/2-1 ; i>=0 ; i--) 277 if(_siftup((PyListObject *)heap, i) == -1) 278 return NULL; 279 Py_INCREF(Py_None); 280 return Py_None; 281} 282 283PyDoc_STRVAR(heapify_doc, 284"Transform list into a heap, in-place, in O(len(heap)) time."); 285 286static PyObject * 287nlargest(PyObject *self, PyObject *args) 288{ 289 PyObject *heap=NULL, *elem, *iterable, *sol, *it, *oldelem; 290 Py_ssize_t i, n; 291 int cmp; 292 293 if (!PyArg_ParseTuple(args, "nO:nlargest", &n, &iterable)) 294 return NULL; 295 296 it = PyObject_GetIter(iterable); 297 if (it == NULL) 298 return NULL; 299 300 heap = PyList_New(0); 301 if (heap == NULL) 302 goto fail; 303 304 for (i=0 ; i<n ; i++ ){ 305 elem = PyIter_Next(it); 306 if (elem == NULL) { 307 if (PyErr_Occurred()) 308 goto fail; 309 else 310 goto sortit; 311 } 312 if (PyList_Append(heap, elem) == -1) { 313 Py_DECREF(elem); 314 goto fail; 315 } 316 Py_DECREF(elem); 317 } 318 if (PyList_GET_SIZE(heap) == 0) 319 goto sortit; 320 321 for (i=n/2-1 ; i>=0 ; i--) 322 if(_siftup((PyListObject *)heap, i) == -1) 323 goto fail; 324 325 sol = PyList_GET_ITEM(heap, 0); 326 while (1) { 327 elem = PyIter_Next(it); 328 if (elem == NULL) { 329 if (PyErr_Occurred()) 330 goto fail; 331 else 332 goto sortit; 333 } 334 cmp = cmp_lt(sol, elem); 335 if (cmp == -1) { 336 Py_DECREF(elem); 337 goto fail; 338 } 339 if (cmp == 0) { 340 Py_DECREF(elem); 341 continue; 342 } 343 oldelem = PyList_GET_ITEM(heap, 0); 344 PyList_SET_ITEM(heap, 0, elem); 345 Py_DECREF(oldelem); 346 if (_siftup((PyListObject *)heap, 0) == -1) 347 goto fail; 348 sol = PyList_GET_ITEM(heap, 0); 349 } 350sortit: 351 if (PyList_Sort(heap) == -1) 352 goto fail; 353 if (PyList_Reverse(heap) == -1) 354 goto fail; 355 Py_DECREF(it); 356 return heap; 357 358fail: 359 Py_DECREF(it); 360 Py_XDECREF(heap); 361 return NULL; 362} 363 364PyDoc_STRVAR(nlargest_doc, 365"Find the n largest elements in a dataset.\n\ 366\n\ 367Equivalent to: sorted(iterable, reverse=True)[:n]\n"); 368 369static int 370_siftdownmax(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos) 371{ 372 PyObject *newitem, *parent; 373 int cmp; 374 Py_ssize_t parentpos; 375 376 assert(PyList_Check(heap)); 377 if (pos >= PyList_GET_SIZE(heap)) { 378 PyErr_SetString(PyExc_IndexError, "index out of range"); 379 return -1; 380 } 381 382 newitem = PyList_GET_ITEM(heap, pos); 383 Py_INCREF(newitem); 384 /* Follow the path to the root, moving parents down until finding 385 a place newitem fits. */ 386 while (pos > startpos){ 387 parentpos = (pos - 1) >> 1; 388 parent = PyList_GET_ITEM(heap, parentpos); 389 cmp = cmp_lt(parent, newitem); 390 if (cmp == -1) { 391 Py_DECREF(newitem); 392 return -1; 393 } 394 if (cmp == 0) 395 break; 396 Py_INCREF(parent); 397 Py_DECREF(PyList_GET_ITEM(heap, pos)); 398 PyList_SET_ITEM(heap, pos, parent); 399 pos = parentpos; 400 } 401 Py_DECREF(PyList_GET_ITEM(heap, pos)); 402 PyList_SET_ITEM(heap, pos, newitem); 403 return 0; 404} 405 406static int 407_siftupmax(PyListObject *heap, Py_ssize_t pos) 408{ 409 Py_ssize_t startpos, endpos, childpos, rightpos; 410 int cmp; 411 PyObject *newitem, *tmp; 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 newitem = PyList_GET_ITEM(heap, pos); 421 Py_INCREF(newitem); 422 423 /* Bubble up the smaller child until hitting a leaf. */ 424 childpos = 2*pos + 1; /* leftmost child position */ 425 while (childpos < endpos) { 426 /* Set childpos to index of smaller child. */ 427 rightpos = childpos + 1; 428 if (rightpos < endpos) { 429 cmp = cmp_lt( 430 PyList_GET_ITEM(heap, rightpos), 431 PyList_GET_ITEM(heap, childpos)); 432 if (cmp == -1) { 433 Py_DECREF(newitem); 434 return -1; 435 } 436 if (cmp == 0) 437 childpos = rightpos; 438 } 439 /* Move the smaller child up. */ 440 tmp = PyList_GET_ITEM(heap, childpos); 441 Py_INCREF(tmp); 442 Py_DECREF(PyList_GET_ITEM(heap, pos)); 443 PyList_SET_ITEM(heap, pos, tmp); 444 pos = childpos; 445 childpos = 2*pos + 1; 446 } 447 448 /* The leaf at pos is empty now. Put newitem there, and and bubble 449 it up to its final resting place (by sifting its parents down). */ 450 Py_DECREF(PyList_GET_ITEM(heap, pos)); 451 PyList_SET_ITEM(heap, pos, newitem); 452 return _siftdownmax(heap, startpos, pos); 453} 454 455static PyObject * 456nsmallest(PyObject *self, PyObject *args) 457{ 458 PyObject *heap=NULL, *elem, *iterable, *los, *it, *oldelem; 459 Py_ssize_t i, n; 460 int cmp; 461 462 if (!PyArg_ParseTuple(args, "nO:nsmallest", &n, &iterable)) 463 return NULL; 464 465 it = PyObject_GetIter(iterable); 466 if (it == NULL) 467 return NULL; 468 469 heap = PyList_New(0); 470 if (heap == NULL) 471 goto fail; 472 473 for (i=0 ; i<n ; i++ ){ 474 elem = PyIter_Next(it); 475 if (elem == NULL) { 476 if (PyErr_Occurred()) 477 goto fail; 478 else 479 goto sortit; 480 } 481 if (PyList_Append(heap, elem) == -1) { 482 Py_DECREF(elem); 483 goto fail; 484 } 485 Py_DECREF(elem); 486 } 487 n = PyList_GET_SIZE(heap); 488 if (n == 0) 489 goto sortit; 490 491 for (i=n/2-1 ; i>=0 ; i--) 492 if(_siftupmax((PyListObject *)heap, i) == -1) 493 goto fail; 494 495 los = PyList_GET_ITEM(heap, 0); 496 while (1) { 497 elem = PyIter_Next(it); 498 if (elem == NULL) { 499 if (PyErr_Occurred()) 500 goto fail; 501 else 502 goto sortit; 503 } 504 cmp = cmp_lt(elem, los); 505 if (cmp == -1) { 506 Py_DECREF(elem); 507 goto fail; 508 } 509 if (cmp == 0) { 510 Py_DECREF(elem); 511 continue; 512 } 513 514 oldelem = PyList_GET_ITEM(heap, 0); 515 PyList_SET_ITEM(heap, 0, elem); 516 Py_DECREF(oldelem); 517 if (_siftupmax((PyListObject *)heap, 0) == -1) 518 goto fail; 519 los = PyList_GET_ITEM(heap, 0); 520 } 521 522sortit: 523 if (PyList_Sort(heap) == -1) 524 goto fail; 525 Py_DECREF(it); 526 return heap; 527 528fail: 529 Py_DECREF(it); 530 Py_XDECREF(heap); 531 return NULL; 532} 533 534PyDoc_STRVAR(nsmallest_doc, 535"Find the n smallest elements in a dataset.\n\ 536\n\ 537Equivalent to: sorted(iterable)[:n]\n"); 538 539static PyMethodDef heapq_methods[] = { 540 {"heappush", (PyCFunction)heappush, 541 METH_VARARGS, heappush_doc}, 542 {"heappushpop", (PyCFunction)heappushpop, 543 METH_VARARGS, heappushpop_doc}, 544 {"heappop", (PyCFunction)heappop, 545 METH_O, heappop_doc}, 546 {"heapreplace", (PyCFunction)heapreplace, 547 METH_VARARGS, heapreplace_doc}, 548 {"heapify", (PyCFunction)heapify, 549 METH_O, heapify_doc}, 550 {"nlargest", (PyCFunction)nlargest, 551 METH_VARARGS, nlargest_doc}, 552 {"nsmallest", (PyCFunction)nsmallest, 553 METH_VARARGS, nsmallest_doc}, 554 {NULL, NULL} /* sentinel */ 555}; 556 557PyDoc_STRVAR(module_doc, 558"Heap queue algorithm (a.k.a. priority queue).\n\ 559\n\ 560Heaps are arrays for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for\n\ 561all k, counting elements from 0. For the sake of comparison,\n\ 562non-existing elements are considered to be infinite. The interesting\n\ 563property of a heap is that a[0] is always its smallest element.\n\ 564\n\ 565Usage:\n\ 566\n\ 567heap = [] # creates an empty heap\n\ 568heappush(heap, item) # pushes a new item on the heap\n\ 569item = heappop(heap) # pops the smallest item from the heap\n\ 570item = heap[0] # smallest item on the heap without popping it\n\ 571heapify(x) # transforms list into a heap, in-place, in linear time\n\ 572item = heapreplace(heap, item) # pops and returns smallest item, and adds\n\ 573 # new item; the heap size is unchanged\n\ 574\n\ 575Our API differs from textbook heap algorithms as follows:\n\ 576\n\ 577- We use 0-based indexing. This makes the relationship between the\n\ 578 index for a node and the indexes for its children slightly less\n\ 579 obvious, but is more suitable since Python uses 0-based indexing.\n\ 580\n\ 581- Our heappop() method returns the smallest item, not the largest.\n\ 582\n\ 583These two make it possible to view the heap as a regular Python list\n\ 584without surprises: heap[0] is the smallest item, and heap.sort()\n\ 585maintains the heap invariant!\n"); 586 587 588PyDoc_STRVAR(__about__, 589"Heap queues\n\ 590\n\ 591[explanation by Fran\xc3\xa7ois Pinard]\n\ 592\n\ 593Heaps are arrays for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for\n\ 594all k, counting elements from 0. For the sake of comparison,\n\ 595non-existing elements are considered to be infinite. The interesting\n\ 596property of a heap is that a[0] is always its smallest element.\n" 597"\n\ 598The strange invariant above is meant to be an efficient memory\n\ 599representation for a tournament. The numbers below are `k', not a[k]:\n\ 600\n\ 601 0\n\ 602\n\ 603 1 2\n\ 604\n\ 605 3 4 5 6\n\ 606\n\ 607 7 8 9 10 11 12 13 14\n\ 608\n\ 609 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30\n\ 610\n\ 611\n\ 612In the tree above, each cell `k' is topping `2*k+1' and `2*k+2'. In\n\ 613an usual binary tournament we see in sports, each cell is the winner\n\ 614over the two cells it tops, and we can trace the winner down the tree\n\ 615to see all opponents s/he had. However, in many computer applications\n\ 616of such tournaments, we do not need to trace the history of a winner.\n\ 617To be more memory efficient, when a winner is promoted, we try to\n\ 618replace it by something else at a lower level, and the rule becomes\n\ 619that a cell and the two cells it tops contain three different items,\n\ 620but the top cell \"wins\" over the two topped cells.\n" 621"\n\ 622If this heap invariant is protected at all time, index 0 is clearly\n\ 623the overall winner. The simplest algorithmic way to remove it and\n\ 624find the \"next\" winner is to move some loser (let's say cell 30 in the\n\ 625diagram above) into the 0 position, and then percolate this new 0 down\n\ 626the tree, exchanging values, until the invariant is re-established.\n\ 627This is clearly logarithmic on the total number of items in the tree.\n\ 628By iterating over all items, you get an O(n ln n) sort.\n" 629"\n\ 630A nice feature of this sort is that you can efficiently insert new\n\ 631items while the sort is going on, provided that the inserted items are\n\ 632not \"better\" than the last 0'th element you extracted. This is\n\ 633especially useful in simulation contexts, where the tree holds all\n\ 634incoming events, and the \"win\" condition means the smallest scheduled\n\ 635time. When an event schedule other events for execution, they are\n\ 636scheduled into the future, so they can easily go into the heap. So, a\n\ 637heap is a good structure for implementing schedulers (this is what I\n\ 638used for my MIDI sequencer :-).\n" 639"\n\ 640Various structures for implementing schedulers have been extensively\n\ 641studied, and heaps are good for this, as they are reasonably speedy,\n\ 642the speed is almost constant, and the worst case is not much different\n\ 643than the average case. However, there are other representations which\n\ 644are more efficient overall, yet the worst cases might be terrible.\n" 645"\n\ 646Heaps are also very useful in big disk sorts. You most probably all\n\ 647know that a big sort implies producing \"runs\" (which are pre-sorted\n\ 648sequences, which size is usually related to the amount of CPU memory),\n\ 649followed by a merging passes for these runs, which merging is often\n\ 650very cleverly organised[1]. It is very important that the initial\n\ 651sort produces the longest runs possible. Tournaments are a good way\n\ 652to that. If, using all the memory available to hold a tournament, you\n\ 653replace and percolate items that happen to fit the current run, you'll\n\ 654produce runs which are twice the size of the memory for random input,\n\ 655and much better for input fuzzily ordered.\n" 656"\n\ 657Moreover, if you output the 0'th item on disk and get an input which\n\ 658may not fit in the current tournament (because the value \"wins\" over\n\ 659the last output value), it cannot fit in the heap, so the size of the\n\ 660heap decreases. The freed memory could be cleverly reused immediately\n\ 661for progressively building a second heap, which grows at exactly the\n\ 662same rate the first heap is melting. When the first heap completely\n\ 663vanishes, you switch heaps and start a new run. Clever and quite\n\ 664effective!\n\ 665\n\ 666In a word, heaps are useful memory structures to know. I use them in\n\ 667a few applications, and I think it is good to keep a `heap' module\n\ 668around. :-)\n" 669"\n\ 670--------------------\n\ 671[1] The disk balancing algorithms which are current, nowadays, are\n\ 672more annoying than clever, and this is a consequence of the seeking\n\ 673capabilities of the disks. On devices which cannot seek, like big\n\ 674tape drives, the story was quite different, and one had to be very\n\ 675clever to ensure (far in advance) that each tape movement will be the\n\ 676most effective possible (that is, will best participate at\n\ 677\"progressing\" the merge). Some tapes were even able to read\n\ 678backwards, and this was also used to avoid the rewinding time.\n\ 679Believe me, real good tape sorts were quite spectacular to watch!\n\ 680From all times, sorting has always been a Great Art! :-)\n"); 681 682 683static struct PyModuleDef _heapqmodule = { 684 PyModuleDef_HEAD_INIT, 685 "_heapq", 686 module_doc, 687 -1, 688 heapq_methods, 689 NULL, 690 NULL, 691 NULL, 692 NULL 693}; 694 695PyMODINIT_FUNC 696PyInit__heapq(void) 697{ 698 PyObject *m, *about; 699 700 m = PyModule_Create(&_heapqmodule); 701 if (m == NULL) 702 return NULL; 703 about = PyUnicode_DecodeUTF8(__about__, strlen(__about__), NULL); 704 PyModule_AddObject(m, "__about__", about); 705 return m; 706} 707 708