summaryrefslogtreecommitdiffstats
path: root/Modules/_heapqmodule.c
diff options
context:
space:
mode:
Diffstat (limited to 'Modules/_heapqmodule.c')
-rw-r--r--Modules/_heapqmodule.c355
1 files changed, 100 insertions, 255 deletions
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index eee56a0..4372ad4 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -9,12 +9,11 @@ annotated by François Pinard, and converted to C by Raymond Hettinger.
#include "Python.h"
static int
-_siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
+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,60 +22,45 @@ _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;
}
static int
-_siftup(PyListObject *heap, Py_ssize_t pos)
+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,38 +73,25 @@ _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);
- return _siftdown(heap, startpos, pos);
+ /* Bubble it up to its final resting place (by sifting its parents down). */
+ return siftdown(heap, startpos, pos);
}
static PyObject *
@@ -139,17 +110,16 @@ heappush(PyObject *self, PyObject *args)
if (PyList_Append(heap, item) == -1)
return NULL;
- if (_siftdown((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-1) == -1)
+ if (siftdown((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-1) == -1)
return NULL;
- Py_INCREF(Py_None);
- return Py_None;
+ Py_RETURN_NONE;
}
PyDoc_STRVAR(heappush_doc,
"heappush(heap, item) -> None. Push item onto heap, maintaining the heap invariant.");
static PyObject *
-heappop(PyObject *self, PyObject *heap)
+heappop_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
{
PyObject *lastelt, *returnitem;
Py_ssize_t n;
@@ -159,7 +129,7 @@ heappop(PyObject *self, PyObject *heap)
return NULL;
}
- /* # raises appropriate IndexError if heap is empty */
+ /* raises IndexError if the heap is empty */
n = PyList_GET_SIZE(heap);
if (n == 0) {
PyErr_SetString(PyExc_IndexError, "index out of range");
@@ -178,18 +148,24 @@ heappop(PyObject *self, PyObject *heap)
return lastelt;
returnitem = PyList_GET_ITEM(heap, 0);
PyList_SET_ITEM(heap, 0, lastelt);
- if (_siftup((PyListObject *)heap, 0) == -1) {
+ if (siftup_func((PyListObject *)heap, 0) == -1) {
Py_DECREF(returnitem);
return NULL;
}
return returnitem;
}
+static PyObject *
+heappop(PyObject *self, PyObject *heap)
+{
+ return heappop_internal(heap, siftup);
+}
+
PyDoc_STRVAR(heappop_doc,
"Pop the smallest item off the heap, maintaining the heap invariant.");
static PyObject *
-heapreplace(PyObject *self, PyObject *args)
+heapreplace_internal(PyObject *args, int siftup_func(PyListObject *, Py_ssize_t))
{
PyObject *heap, *item, *returnitem;
@@ -209,13 +185,19 @@ heapreplace(PyObject *self, PyObject *args)
returnitem = PyList_GET_ITEM(heap, 0);
Py_INCREF(item);
PyList_SET_ITEM(heap, 0, item);
- if (_siftup((PyListObject *)heap, 0) == -1) {
+ if (siftup_func((PyListObject *)heap, 0) == -1) {
Py_DECREF(returnitem);
return NULL;
}
return returnitem;
}
+static PyObject *
+heapreplace(PyObject *self, PyObject *args)
+{
+ return heapreplace_internal(args, siftup);
+}
+
PyDoc_STRVAR(heapreplace_doc,
"heapreplace(heap, item) -> value. Pop and return the current smallest value, and add the new item.\n\
\n\
@@ -256,7 +238,7 @@ heappushpop(PyObject *self, PyObject *args)
returnitem = PyList_GET_ITEM(heap, 0);
Py_INCREF(item);
PyList_SET_ITEM(heap, 0, item);
- if (_siftup((PyListObject *)heap, 0) == -1) {
+ if (siftup((PyListObject *)heap, 0) == -1) {
Py_DECREF(returnitem);
return NULL;
}
@@ -269,7 +251,7 @@ from the heap. The combined action runs more efficiently than\n\
heappush() followed by a separate call to heappop().");
static PyObject *
-heapify(PyObject *self, PyObject *heap)
+heapify_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
{
Py_ssize_t i, n;
@@ -287,141 +269,65 @@ heapify(PyObject *self, PyObject *heap)
and that's again n//2-1.
*/
for (i=n/2-1 ; i>=0 ; i--)
- if(_siftup((PyListObject *)heap, i) == -1)
+ if(siftup_func((PyListObject *)heap, i) == -1)
return NULL;
- Py_INCREF(Py_None);
- return Py_None;
+ Py_RETURN_NONE;
}
-PyDoc_STRVAR(heapify_doc,
-"Transform list into a heap, in-place, in O(len(heap)) time.");
-
static PyObject *
-nlargest(PyObject *self, PyObject *args)
+heapify(PyObject *self, PyObject *heap)
{
- 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;
+ return heapify_internal(heap, siftup);
}
-PyDoc_STRVAR(nlargest_doc,
-"Find the n largest elements in a dataset.\n\
-\n\
-Equivalent to: sorted(iterable, reverse=True)[:n]\n");
+PyDoc_STRVAR(heapify_doc,
+"Transform list into a heap, in-place, in O(len(heap)) time.");
static int
-_siftdownmax(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
+siftdown_max(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;
}
static int
-_siftupmax(PyListObject *heap, Py_ssize_t pos)
+siftup_max(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 +336,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 +348,50 @@ _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);
- return _siftdownmax(heap, startpos, pos);
+ /* Bubble it up to its final resting place (by sifting its parents down). */
+ return siftdown_max(heap, startpos, pos);
}
static PyObject *
-nsmallest(PyObject *self, PyObject *args)
+heappop_max(PyObject *self, PyObject *heap)
{
- PyObject *heap=NULL, *elem, *iterable, *los, *it, *oldelem;
- Py_ssize_t i, n;
- int cmp;
-
- if (!PyArg_ParseTuple(args, "nO:nsmallest", &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);
- }
- n = PyList_GET_SIZE(heap);
- if (n == 0)
- goto sortit;
+ return heappop_internal(heap, siftup_max);
+}
- 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;
- }
+PyDoc_STRVAR(heappop_max_doc, "Maxheap variant of heappop.");
- 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);
- }
+static PyObject *
+heapreplace_max(PyObject *self, PyObject *args)
+{
+ return heapreplace_internal(args, siftup_max);
+}
-sortit:
- if (PyList_Sort(heap) == -1)
- goto fail;
- Py_DECREF(it);
- return heap;
+PyDoc_STRVAR(heapreplace_max_doc, "Maxheap variant of heapreplace");
-fail:
- Py_DECREF(it);
- Py_XDECREF(heap);
- return NULL;
+static PyObject *
+heapify_max(PyObject *self, PyObject *heap)
+{
+ return heapify_internal(heap, siftup_max);
}
-PyDoc_STRVAR(nsmallest_doc,
-"Find the n smallest elements in a dataset.\n\
-\n\
-Equivalent to: sorted(iterable)[:n]\n");
+PyDoc_STRVAR(heapify_max_doc, "Maxheap variant of heapify.");
static PyMethodDef heapq_methods[] = {
{"heappush", (PyCFunction)heappush,
@@ -561,10 +404,12 @@ 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},
+ {"_heappop_max", (PyCFunction)heappop_max,
+ METH_O, heappop_max_doc},
+ {"_heapreplace_max",(PyCFunction)heapreplace_max,
+ METH_VARARGS, heapreplace_max_doc},
+ {"_heapify_max", (PyCFunction)heapify_max,
+ METH_O, heapify_max_doc},
{NULL, NULL} /* sentinel */
};
generated by cgit v0.12 at 2025-01-07 20:12:20 (GMT)