diff options
Diffstat (limited to 'Modules/_heapqmodule.c')
| -rw-r--r-- | Modules/_heapqmodule.c | 432 |
1 files changed, 191 insertions, 241 deletions
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c index 04845f165a..c343862b8c 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 */ }; |