summaryrefslogtreecommitdiffstats
path: root/Modules/_heapqmodule.c
diff options
context:
space:
mode:
Diffstat (limited to 'Modules/_heapqmodule.c')
-rw-r--r--Modules/_heapqmodule.c434
1 files changed, 192 insertions, 242 deletions
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 04845f1..136abf5 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 */
};
@@ -600,7 +550,7 @@ representation for a tournament. The numbers below are `k', not a[k]:\n\
\n\
\n\
In the tree above, each cell `k' is topping `2*k+1' and `2*k+2'. In\n\
-an usual binary tournament we see in sports, each cell is the winner\n\
+a usual binary tournament we see in sports, each cell is the winner\n\
over the two cells it tops, and we can trace the winner down the tree\n\
to see all opponents s/he had. However, in many computer applications\n\
of such tournaments, we do not need to trace the history of a winner.\n\
generated by cgit v0.12 at 2025-12-02 17:05:24 (GMT)