summaryrefslogtreecommitdiffstats
path: root/Modules/_heapqmodule.c
diff options
context:
space:
mode:
Diffstat (limited to 'Modules/_heapqmodule.c')
-rw-r--r--Modules/_heapqmodule.c432
1 files changed, 191 insertions, 241 deletions
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 04845f1..c343862 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -9,9 +9,9 @@ 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;
+ PyObject *newitem, *parent, **arr;
Py_ssize_t parentpos, size;
int cmp;
@@ -24,12 +24,13 @@ _siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
/* Follow the path to the root, moving parents down until finding
a place newitem fits. */
- newitem = PyList_GET_ITEM(heap, pos);
+ arr = _PyList_ITEMS(heap);
+ newitem = arr[pos];
while (pos > startpos) {
parentpos = (pos - 1) >> 1;
- parent = PyList_GET_ITEM(heap, parentpos);
+ parent = arr[parentpos];
cmp = PyObject_RichCompareBool(newitem, parent, Py_LT);
- if (cmp == -1)
+ if (cmp < 0)
return -1;
if (size != PyList_GET_SIZE(heap)) {
PyErr_SetString(PyExc_RuntimeError,
@@ -38,20 +39,21 @@ _siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
}
if (cmp == 0)
break;
- parent = PyList_GET_ITEM(heap, parentpos);
- newitem = PyList_GET_ITEM(heap, pos);
- PyList_SET_ITEM(heap, parentpos, newitem);
- PyList_SET_ITEM(heap, pos, parent);
+ arr = _PyList_ITEMS(heap);
+ parent = arr[parentpos];
+ newitem = arr[pos];
+ arr[parentpos] = newitem;
+ arr[pos] = parent;
pos = parentpos;
}
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;
+ Py_ssize_t startpos, endpos, childpos, limit;
+ PyObject *tmp1, *tmp2, **arr;
int cmp;
assert(PyList_Check(heap));
@@ -63,20 +65,19 @@ _siftup(PyListObject *heap, Py_ssize_t pos)
}
/* Bubble up the smaller child until hitting a leaf. */
+ arr = _PyList_ITEMS(heap);
limit = endpos / 2; /* smallest pos that has no child */
while (pos < limit) {
/* Set childpos to index of smaller child. */
childpos = 2*pos + 1; /* leftmost child position */
- rightpos = childpos + 1;
- if (rightpos < endpos) {
+ if (childpos + 1 < endpos) {
cmp = PyObject_RichCompareBool(
- PyList_GET_ITEM(heap, childpos),
- PyList_GET_ITEM(heap, rightpos),
+ arr[childpos],
+ arr[childpos + 1],
Py_LT);
- if (cmp == -1)
+ if (cmp < 0)
return -1;
- if (cmp == 0)
- childpos = rightpos;
+ childpos += ((unsigned)cmp ^ 1); /* increment when cmp==0 */
if (endpos != PyList_GET_SIZE(heap)) {
PyErr_SetString(PyExc_RuntimeError,
"list changed size during iteration");
@@ -84,14 +85,15 @@ _siftup(PyListObject *heap, Py_ssize_t pos)
}
}
/* Move the smaller child up. */
- tmp1 = PyList_GET_ITEM(heap, childpos);
- tmp2 = PyList_GET_ITEM(heap, pos);
- PyList_SET_ITEM(heap, childpos, tmp2);
- PyList_SET_ITEM(heap, pos, tmp1);
+ arr = _PyList_ITEMS(heap);
+ tmp1 = arr[childpos];
+ tmp2 = arr[pos];
+ arr[childpos] = tmp2;
+ arr[pos] = tmp1;
pos = childpos;
}
/* Bubble it up to its final resting place (by sifting its parents down). */
- return _siftdown(heap, startpos, pos);
+ return siftdown(heap, startpos, pos);
}
static PyObject *
@@ -107,20 +109,19 @@ heappush(PyObject *self, PyObject *args)
return NULL;
}
- if (PyList_Append(heap, item) == -1)
+ if (PyList_Append(heap, item))
return NULL;
- if (_siftdown((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-1) == -1)
+ if (siftdown((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-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;
@@ -130,7 +131,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");
@@ -139,7 +140,7 @@ heappop(PyObject *self, PyObject *heap)
lastelt = PyList_GET_ITEM(heap, n-1) ;
Py_INCREF(lastelt);
- if (PyList_SetSlice(heap, n-1, n, NULL) < 0) {
+ if (PyList_SetSlice(heap, n-1, n, NULL)) {
Py_DECREF(lastelt);
return NULL;
}
@@ -149,18 +150,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)) {
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;
@@ -172,7 +179,7 @@ heapreplace(PyObject *self, PyObject *args)
return NULL;
}
- if (PyList_GET_SIZE(heap) < 1) {
+ if (PyList_GET_SIZE(heap) == 0) {
PyErr_SetString(PyExc_IndexError, "index out of range");
return NULL;
}
@@ -180,13 +187,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)) {
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\
@@ -211,13 +224,13 @@ heappushpop(PyObject *self, PyObject *args)
return NULL;
}
- if (PyList_GET_SIZE(heap) < 1) {
+ if (PyList_GET_SIZE(heap) == 0) {
Py_INCREF(item);
return item;
}
cmp = PyObject_RichCompareBool(PyList_GET_ITEM(heap, 0), item, Py_LT);
- if (cmp == -1)
+ if (cmp < 0)
return NULL;
if (cmp == 0) {
Py_INCREF(item);
@@ -232,7 +245,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)) {
Py_DECREF(returnitem);
return NULL;
}
@@ -244,8 +257,73 @@ PyDoc_STRVAR(heappushpop_doc,
from the heap. The combined action runs more efficiently than\n\
heappush() followed by a separate call to heappop().");
+static Py_ssize_t
+keep_top_bit(Py_ssize_t n)
+{
+ int i = 0;
+
+ while (n > 1) {
+ n >>= 1;
+ i++;
+ }
+ return n << i;
+}
+
+/* Cache friendly version of heapify()
+ -----------------------------------
+
+ Build-up a heap in O(n) time by performing siftup() operations
+ on nodes whose children are already heaps.
+
+ The simplest way is to sift the nodes in reverse order from
+ n//2-1 to 0 inclusive. The downside is that children may be
+ out of cache by the time their parent is reached.
+
+ A better way is to not wait for the children to go out of cache.
+ Once a sibling pair of child nodes have been sifted, immediately
+ sift their parent node (while the children are still in cache).
+
+ Both ways build child heaps before their parents, so both ways
+ do the exact same number of comparisons and produce exactly
+ the same heap. The only difference is that the traversal
+ order is optimized for cache efficiency.
+*/
+
static PyObject *
-heapify(PyObject *self, PyObject *heap)
+cache_friendly_heapify(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
+{
+ Py_ssize_t i, j, m, mhalf, leftmost;
+
+ m = PyList_GET_SIZE(heap) >> 1; /* index of first childless node */
+ leftmost = keep_top_bit(m + 1) - 1; /* leftmost node in row of m */
+ mhalf = m >> 1; /* parent of first childless node */
+
+ for (i = leftmost - 1 ; i >= mhalf ; i--) {
+ j = i;
+ while (1) {
+ if (siftup_func((PyListObject *)heap, j))
+ return NULL;
+ if (!(j & 1))
+ break;
+ j >>= 1;
+ }
+ }
+
+ for (i = m - 1 ; i >= leftmost ; i--) {
+ j = i;
+ while (1) {
+ if (siftup_func((PyListObject *)heap, j))
+ return NULL;
+ if (!(j & 1))
+ break;
+ j >>= 1;
+ }
+ }
+ Py_RETURN_NONE;
+}
+
+static PyObject *
+heapify_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
{
Py_ssize_t i, n;
@@ -254,7 +332,14 @@ heapify(PyObject *self, PyObject *heap)
return NULL;
}
+ /* For heaps likely to be bigger than L1 cache, we use the cache
+ friendly heapify function. For smaller heaps that fit entirely
+ in cache, we prefer the simpler algorithm with less branching.
+ */
n = PyList_GET_SIZE(heap);
+ if (n > 2500)
+ return cache_friendly_heapify(heap, siftup_func);
+
/* Transform bottom-up. The largest index there's any point to
looking at is the largest with a child index in-range, so must
have 2*i + 1 < n, or i < (n-1)/2. If n is even = 2*j, this is
@@ -262,142 +347,68 @@ heapify(PyObject *self, PyObject *heap)
n is odd = 2*j+1, this is (2*j+1-1)/2 = j so j-1 is the largest,
and that's again n//2-1.
*/
- for (i=n/2-1 ; i>=0 ; i--)
- if(_siftup((PyListObject *)heap, i) == -1)
+ for (i = n/2 - 1 ; i >= 0 ; i--)
+ if (siftup_func((PyListObject *)heap, i))
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;
+ PyObject *newitem, *parent, **arr;
+ 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){
+ arr = _PyList_ITEMS(heap);
+ newitem = arr[pos];
+ while (pos > startpos) {
parentpos = (pos - 1) >> 1;
- parent = PyList_GET_ITEM(heap, parentpos);
+ parent = arr[parentpos];
cmp = PyObject_RichCompareBool(parent, newitem, Py_LT);
- if (cmp == -1) {
- Py_DECREF(newitem);
+ if (cmp < 0)
+ 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));
- PyList_SET_ITEM(heap, pos, parent);
+ arr = _PyList_ITEMS(heap);
+ parent = arr[parentpos];
+ newitem = arr[pos];
+ arr[parentpos] = newitem;
+ arr[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;
+ Py_ssize_t startpos, endpos, childpos, limit;
+ PyObject *tmp1, *tmp2, **arr;
int cmp;
- PyObject *newitem, *tmp;
assert(PyList_Check(heap));
endpos = PyList_GET_SIZE(heap);
@@ -406,125 +417,62 @@ _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. */
+ arr = _PyList_ITEMS(heap);
limit = endpos / 2; /* smallest pos that has no child */
while (pos < limit) {
/* Set childpos to index of smaller child. */
childpos = 2*pos + 1; /* leftmost child position */
- rightpos = childpos + 1;
- if (rightpos < endpos) {
+ if (childpos + 1 < endpos) {
cmp = PyObject_RichCompareBool(
- PyList_GET_ITEM(heap, rightpos),
- PyList_GET_ITEM(heap, childpos),
+ arr[childpos + 1],
+ arr[childpos],
Py_LT);
- if (cmp == -1) {
- Py_DECREF(newitem);
+ if (cmp < 0)
+ return -1;
+ childpos += ((unsigned)cmp ^ 1); /* increment when cmp==0 */
+ if (endpos != PyList_GET_SIZE(heap)) {
+ PyErr_SetString(PyExc_RuntimeError,
+ "list changed size during iteration");
return -1;
}
- if (cmp == 0)
- childpos = rightpos;
}
/* 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);
+ arr = _PyList_ITEMS(heap);
+ tmp1 = arr[childpos];
+ tmp2 = arr[pos];
+ arr[childpos] = tmp2;
+ arr[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;
+ return heappop_internal(heap, siftup_max);
+}
- it = PyObject_GetIter(iterable);
- if (it == NULL)
- return NULL;
+PyDoc_STRVAR(heappop_max_doc, "Maxheap variant of heappop.");
- 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);
- }
+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,
@@ -537,10 +485,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 2024-12-28 04:40:30 (GMT)