summaryrefslogtreecommitdiffstats
path: root/Modules/_heapqmodule.c
diff options
context:
space:
mode:
Diffstat (limited to 'Modules/_heapqmodule.c')
-rw-r--r--Modules/_heapqmodule.c299
1 files changed, 65 insertions, 234 deletions
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index eee56a0..ad190df 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -11,10 +11,9 @@ annotated by François Pinard, and converted to C by Raymond Hettinger.
static int
_siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
{
- PyObject *newitem, *parent, *olditem;
+ PyObject *newitem, *parent;
+ Py_ssize_t parentpos, size;
int cmp;
- Py_ssize_t parentpos;
- Py_ssize_t size;
assert(PyList_Check(heap));
size = PyList_GET_SIZE(heap);
@@ -23,39 +22,28 @@ _siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
return -1;
}
- newitem = PyList_GET_ITEM(heap, pos);
- Py_INCREF(newitem);
/* Follow the path to the root, moving parents down until finding
a place newitem fits. */
- while (pos > startpos){
+ newitem = PyList_GET_ITEM(heap, pos);
+ while (pos > startpos) {
parentpos = (pos - 1) >> 1;
parent = PyList_GET_ITEM(heap, parentpos);
cmp = PyObject_RichCompareBool(newitem, parent, Py_LT);
- if (cmp == -1) {
- Py_DECREF(newitem);
+ if (cmp == -1)
return -1;
- }
if (size != PyList_GET_SIZE(heap)) {
- Py_DECREF(newitem);
PyErr_SetString(PyExc_RuntimeError,
"list changed size during iteration");
return -1;
}
if (cmp == 0)
break;
- Py_INCREF(parent);
- olditem = PyList_GET_ITEM(heap, pos);
+ parent = PyList_GET_ITEM(heap, parentpos);
+ newitem = PyList_GET_ITEM(heap, pos);
+ PyList_SET_ITEM(heap, parentpos, newitem);
PyList_SET_ITEM(heap, pos, parent);
- Py_DECREF(olditem);
pos = parentpos;
- if (size != PyList_GET_SIZE(heap)) {
- PyErr_SetString(PyExc_RuntimeError,
- "list changed size during iteration");
- return -1;
- }
}
- Py_DECREF(PyList_GET_ITEM(heap, pos));
- PyList_SET_ITEM(heap, pos, newitem);
return 0;
}
@@ -63,20 +51,16 @@ static int
_siftup(PyListObject *heap, Py_ssize_t pos)
{
Py_ssize_t startpos, endpos, childpos, rightpos, limit;
+ PyObject *tmp1, *tmp2;
int cmp;
- PyObject *newitem, *tmp, *olditem;
- Py_ssize_t size;
assert(PyList_Check(heap));
- size = PyList_GET_SIZE(heap);
- endpos = size;
+ endpos = PyList_GET_SIZE(heap);
startpos = pos;
if (pos >= endpos) {
PyErr_SetString(PyExc_IndexError, "index out of range");
return -1;
}
- newitem = PyList_GET_ITEM(heap, pos);
- Py_INCREF(newitem);
/* Bubble up the smaller child until hitting a leaf. */
limit = endpos / 2; /* smallest pos that has no child */
@@ -89,37 +73,24 @@ _siftup(PyListObject *heap, Py_ssize_t pos)
PyList_GET_ITEM(heap, childpos),
PyList_GET_ITEM(heap, rightpos),
Py_LT);
- if (cmp == -1) {
- Py_DECREF(newitem);
+ if (cmp == -1)
return -1;
- }
if (cmp == 0)
childpos = rightpos;
- }
- if (size != PyList_GET_SIZE(heap)) {
- Py_DECREF(newitem);
- PyErr_SetString(PyExc_RuntimeError,
- "list changed size during iteration");
- return -1;
+ if (endpos != PyList_GET_SIZE(heap)) {
+ PyErr_SetString(PyExc_RuntimeError,
+ "list changed size during iteration");
+ return -1;
+ }
}
/* Move the smaller child up. */
- tmp = PyList_GET_ITEM(heap, childpos);
- Py_INCREF(tmp);
- olditem = PyList_GET_ITEM(heap, pos);
- PyList_SET_ITEM(heap, pos, tmp);
- Py_DECREF(olditem);
+ tmp1 = PyList_GET_ITEM(heap, childpos);
+ tmp2 = PyList_GET_ITEM(heap, pos);
+ PyList_SET_ITEM(heap, childpos, tmp2);
+ PyList_SET_ITEM(heap, pos, tmp1);
pos = childpos;
- if (size != PyList_GET_SIZE(heap)) {
- PyErr_SetString(PyExc_RuntimeError,
- "list changed size during iteration");
- return -1;
- }
}
-
- /* The leaf at pos is empty now. Put newitem there, and bubble
- it up to its final resting place (by sifting its parents down). */
- Py_DECREF(PyList_GET_ITEM(heap, pos));
- PyList_SET_ITEM(heap, pos, newitem);
+ /* Bubble it up to its final resting place (by sifting its parents down). */
return _siftdown(heap, startpos, pos);
}
@@ -296,123 +267,42 @@ heapify(PyObject *self, PyObject *heap)
PyDoc_STRVAR(heapify_doc,
"Transform list into a heap, in-place, in O(len(heap)) time.");
-static PyObject *
-nlargest(PyObject *self, PyObject *args)
-{
- PyObject *heap=NULL, *elem, *iterable, *sol, *it, *oldelem;
- Py_ssize_t i, n;
- int cmp;
-
- if (!PyArg_ParseTuple(args, "nO:nlargest", &n, &iterable))
- return NULL;
-
- it = PyObject_GetIter(iterable);
- if (it == NULL)
- return NULL;
-
- heap = PyList_New(0);
- if (heap == NULL)
- goto fail;
-
- for (i=0 ; i<n ; i++ ){
- elem = PyIter_Next(it);
- if (elem == NULL) {
- if (PyErr_Occurred())
- goto fail;
- else
- goto sortit;
- }
- if (PyList_Append(heap, elem) == -1) {
- Py_DECREF(elem);
- goto fail;
- }
- Py_DECREF(elem);
- }
- if (PyList_GET_SIZE(heap) == 0)
- goto sortit;
-
- for (i=n/2-1 ; i>=0 ; i--)
- if(_siftup((PyListObject *)heap, i) == -1)
- goto fail;
-
- sol = PyList_GET_ITEM(heap, 0);
- while (1) {
- elem = PyIter_Next(it);
- if (elem == NULL) {
- if (PyErr_Occurred())
- goto fail;
- else
- goto sortit;
- }
- cmp = PyObject_RichCompareBool(sol, elem, Py_LT);
- if (cmp == -1) {
- Py_DECREF(elem);
- goto fail;
- }
- if (cmp == 0) {
- Py_DECREF(elem);
- continue;
- }
- oldelem = PyList_GET_ITEM(heap, 0);
- PyList_SET_ITEM(heap, 0, elem);
- Py_DECREF(oldelem);
- if (_siftup((PyListObject *)heap, 0) == -1)
- goto fail;
- sol = PyList_GET_ITEM(heap, 0);
- }
-sortit:
- if (PyList_Sort(heap) == -1)
- goto fail;
- if (PyList_Reverse(heap) == -1)
- goto fail;
- Py_DECREF(it);
- return heap;
-
-fail:
- Py_DECREF(it);
- Py_XDECREF(heap);
- return NULL;
-}
-
-PyDoc_STRVAR(nlargest_doc,
-"Find the n largest elements in a dataset.\n\
-\n\
-Equivalent to: sorted(iterable, reverse=True)[:n]\n");
-
static int
_siftdownmax(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
{
PyObject *newitem, *parent;
+ Py_ssize_t parentpos, size;
int cmp;
- Py_ssize_t parentpos;
assert(PyList_Check(heap));
- if (pos >= PyList_GET_SIZE(heap)) {
+ size = PyList_GET_SIZE(heap);
+ if (pos >= size) {
PyErr_SetString(PyExc_IndexError, "index out of range");
return -1;
}
- newitem = PyList_GET_ITEM(heap, pos);
- Py_INCREF(newitem);
/* Follow the path to the root, moving parents down until finding
a place newitem fits. */
- while (pos > startpos){
+ newitem = PyList_GET_ITEM(heap, pos);
+ while (pos > startpos) {
parentpos = (pos - 1) >> 1;
parent = PyList_GET_ITEM(heap, parentpos);
cmp = PyObject_RichCompareBool(parent, newitem, Py_LT);
- if (cmp == -1) {
- Py_DECREF(newitem);
+ if (cmp == -1)
+ return -1;
+ if (size != PyList_GET_SIZE(heap)) {
+ PyErr_SetString(PyExc_RuntimeError,
+ "list changed size during iteration");
return -1;
}
if (cmp == 0)
break;
- Py_INCREF(parent);
- Py_DECREF(PyList_GET_ITEM(heap, pos));
+ parent = PyList_GET_ITEM(heap, parentpos);
+ newitem = PyList_GET_ITEM(heap, pos);
+ PyList_SET_ITEM(heap, parentpos, newitem);
PyList_SET_ITEM(heap, pos, parent);
pos = parentpos;
}
- Py_DECREF(PyList_GET_ITEM(heap, pos));
- PyList_SET_ITEM(heap, pos, newitem);
return 0;
}
@@ -420,8 +310,8 @@ static int
_siftupmax(PyListObject *heap, Py_ssize_t pos)
{
Py_ssize_t startpos, endpos, childpos, rightpos, limit;
+ PyObject *tmp1, *tmp2;
int cmp;
- PyObject *newitem, *tmp;
assert(PyList_Check(heap));
endpos = PyList_GET_SIZE(heap);
@@ -430,8 +320,6 @@ _siftupmax(PyListObject *heap, Py_ssize_t pos)
PyErr_SetString(PyExc_IndexError, "index out of range");
return -1;
}
- newitem = PyList_GET_ITEM(heap, pos);
- Py_INCREF(newitem);
/* Bubble up the smaller child until hitting a leaf. */
limit = endpos / 2; /* smallest pos that has no child */
@@ -444,111 +332,56 @@ _siftupmax(PyListObject *heap, Py_ssize_t pos)
PyList_GET_ITEM(heap, rightpos),
PyList_GET_ITEM(heap, childpos),
Py_LT);
- if (cmp == -1) {
- Py_DECREF(newitem);
+ if (cmp == -1)
return -1;
- }
if (cmp == 0)
childpos = rightpos;
+ if (endpos != PyList_GET_SIZE(heap)) {
+ PyErr_SetString(PyExc_RuntimeError,
+ "list changed size during iteration");
+ return -1;
+ }
}
/* Move the smaller child up. */
- tmp = PyList_GET_ITEM(heap, childpos);
- Py_INCREF(tmp);
- Py_DECREF(PyList_GET_ITEM(heap, pos));
- PyList_SET_ITEM(heap, pos, tmp);
+ tmp1 = PyList_GET_ITEM(heap, childpos);
+ tmp2 = PyList_GET_ITEM(heap, pos);
+ PyList_SET_ITEM(heap, childpos, tmp2);
+ PyList_SET_ITEM(heap, pos, tmp1);
pos = childpos;
}
-
- /* The leaf at pos is empty now. Put newitem there, and bubble
- it up to its final resting place (by sifting its parents down). */
- Py_DECREF(PyList_GET_ITEM(heap, pos));
- PyList_SET_ITEM(heap, pos, newitem);
+ /* Bubble it up to its final resting place (by sifting its parents down). */
return _siftdownmax(heap, startpos, pos);
}
static PyObject *
-nsmallest(PyObject *self, PyObject *args)
+_heapreplace_max(PyObject *self, PyObject *args)
{
- PyObject *heap=NULL, *elem, *iterable, *los, *it, *oldelem;
- Py_ssize_t i, n;
- int cmp;
+ PyObject *heap, *item, *returnitem;
- if (!PyArg_ParseTuple(args, "nO:nsmallest", &n, &iterable))
+ if (!PyArg_UnpackTuple(args, "_heapreplace_max", 2, 2, &heap, &item))
return NULL;
- it = PyObject_GetIter(iterable);
- if (it == NULL)
+ if (!PyList_Check(heap)) {
+ PyErr_SetString(PyExc_TypeError, "heap argument must be a list");
return NULL;
-
- heap = PyList_New(0);
- if (heap == NULL)
- goto fail;
-
- for (i=0 ; i<n ; i++ ){
- elem = PyIter_Next(it);
- if (elem == NULL) {
- if (PyErr_Occurred())
- goto fail;
- else
- goto sortit;
- }
- if (PyList_Append(heap, elem) == -1) {
- Py_DECREF(elem);
- goto fail;
- }
- Py_DECREF(elem);
}
- n = PyList_GET_SIZE(heap);
- if (n == 0)
- goto sortit;
- for (i=n/2-1 ; i>=0 ; i--)
- if(_siftupmax((PyListObject *)heap, i) == -1)
- goto fail;
-
- los = PyList_GET_ITEM(heap, 0);
- while (1) {
- elem = PyIter_Next(it);
- if (elem == NULL) {
- if (PyErr_Occurred())
- goto fail;
- else
- goto sortit;
- }
- cmp = PyObject_RichCompareBool(elem, los, Py_LT);
- if (cmp == -1) {
- Py_DECREF(elem);
- goto fail;
- }
- if (cmp == 0) {
- Py_DECREF(elem);
- continue;
- }
-
- oldelem = PyList_GET_ITEM(heap, 0);
- PyList_SET_ITEM(heap, 0, elem);
- Py_DECREF(oldelem);
- if (_siftupmax((PyListObject *)heap, 0) == -1)
- goto fail;
- los = PyList_GET_ITEM(heap, 0);
+ if (PyList_GET_SIZE(heap) < 1) {
+ PyErr_SetString(PyExc_IndexError, "index out of range");
+ return NULL;
}
-sortit:
- if (PyList_Sort(heap) == -1)
- goto fail;
- Py_DECREF(it);
- return heap;
-
-fail:
- Py_DECREF(it);
- Py_XDECREF(heap);
- return NULL;
+ returnitem = PyList_GET_ITEM(heap, 0);
+ Py_INCREF(item);
+ PyList_SET_ITEM(heap, 0, item);
+ if (_siftupmax((PyListObject *)heap, 0) == -1) {
+ Py_DECREF(returnitem);
+ return NULL;
+ }
+ return returnitem;
}
-PyDoc_STRVAR(nsmallest_doc,
-"Find the n smallest elements in a dataset.\n\
-\n\
-Equivalent to: sorted(iterable)[:n]\n");
+PyDoc_STRVAR(heapreplace_max_doc, "Maxheap variant of heapreplace");
static PyMethodDef heapq_methods[] = {
{"heappush", (PyCFunction)heappush,
@@ -561,10 +394,8 @@ static PyMethodDef heapq_methods[] = {
METH_VARARGS, heapreplace_doc},
{"heapify", (PyCFunction)heapify,
METH_O, heapify_doc},
- {"nlargest", (PyCFunction)nlargest,
- METH_VARARGS, nlargest_doc},
- {"nsmallest", (PyCFunction)nsmallest,
- METH_VARARGS, nsmallest_doc},
+ {"_heapreplace_max",(PyCFunction)_heapreplace_max,
+ METH_VARARGS, heapreplace_max_doc},
{NULL, NULL} /* sentinel */
};
generated by cgit v0.12 at 2025-11-27 17:49:23 (GMT)