diff options
Diffstat (limited to 'Modules/_heapqmodule.c')
| -rw-r--r-- | Modules/_heapqmodule.c | 97 |
1 files changed, 38 insertions, 59 deletions
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c index f377e9c..04845f1 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,103 +22,75 @@ _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) { - Py_ssize_t startpos, endpos, childpos, rightpos; + 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. */ - childpos = 2*pos + 1; /* leftmost child position */ - while (childpos < endpos) { + 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) { cmp = PyObject_RichCompareBool( 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; - childpos = 2*pos + 1; - 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); } @@ -168,7 +139,10 @@ heappop(PyObject *self, PyObject *heap) lastelt = PyList_GET_ITEM(heap, n-1) ; Py_INCREF(lastelt); - PyList_SetSlice(heap, n-1, n, NULL); + if (PyList_SetSlice(heap, n-1, n, NULL) < 0) { + Py_DECREF(lastelt); + return NULL; + } n--; if (!n) @@ -250,6 +224,11 @@ heappushpop(PyObject *self, PyObject *args) return item; } + if (PyList_GET_SIZE(heap) == 0) { + PyErr_SetString(PyExc_IndexError, "index out of range"); + return NULL; + } + returnitem = PyList_GET_ITEM(heap, 0); Py_INCREF(item); PyList_SET_ITEM(heap, 0, item); @@ -416,7 +395,7 @@ _siftdownmax(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos) static int _siftupmax(PyListObject *heap, Py_ssize_t pos) { - Py_ssize_t startpos, endpos, childpos, rightpos; + Py_ssize_t startpos, endpos, childpos, rightpos, limit; int cmp; PyObject *newitem, *tmp; @@ -431,9 +410,10 @@ _siftupmax(PyListObject *heap, Py_ssize_t pos) Py_INCREF(newitem); /* Bubble up the smaller child until hitting a leaf. */ - childpos = 2*pos + 1; /* leftmost child position */ - while (childpos < endpos) { + 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) { cmp = PyObject_RichCompareBool( @@ -453,7 +433,6 @@ _siftupmax(PyListObject *heap, Py_ssize_t pos) Py_DECREF(PyList_GET_ITEM(heap, pos)); PyList_SET_ITEM(heap, pos, tmp); pos = childpos; - childpos = 2*pos + 1; } /* The leaf at pos is empty now. Put newitem there, and bubble |