Issue 2274: Add heapq.heappushpop().

This commit is contained in:
Raymond Hettinger 2008-03-13 19:03:51 +00:00
parent 431f029486
commit 53bdf09343
5 changed files with 95 additions and 2 deletions

View file

@ -162,6 +162,11 @@ heapreplace(PyObject *self, PyObject *args)
{
PyObject *heap, *item, *returnitem;
if (Py_Py3kWarningFlag &&
PyErr_Warn(PyExc_DeprecationWarning,
"In 3.x, heapreplace() was removed. Use heappushpop() instead.") < 0)
return NULL;
if (!PyArg_UnpackTuple(args, "heapreplace", 2, 2, &heap, &item))
return NULL;
@ -195,6 +200,48 @@ this routine unless written as part of a conditional replacement:\n\n\
if item > heap[0]:\n\
item = heapreplace(heap, item)\n");
static PyObject *
heappushpop(PyObject *self, PyObject *args)
{
PyObject *heap, *item, *returnitem;
int cmp;
if (!PyArg_UnpackTuple(args, "heappushpop", 2, 2, &heap, &item))
return NULL;
if (!PyList_Check(heap)) {
PyErr_SetString(PyExc_TypeError, "heap argument must be a list");
return NULL;
}
if (PyList_GET_SIZE(heap) < 1) {
Py_INCREF(item);
return item;
}
cmp = PyObject_RichCompareBool(item, PyList_GET_ITEM(heap, 0), Py_LE);
if (cmp == -1)
return NULL;
if (cmp == 1) {
Py_INCREF(item);
return item;
}
returnitem = PyList_GET_ITEM(heap, 0);
Py_INCREF(item);
PyList_SET_ITEM(heap, 0, item);
if (_siftup((PyListObject *)heap, 0) == -1) {
Py_DECREF(returnitem);
return NULL;
}
return returnitem;
}
PyDoc_STRVAR(heappushpop_doc,
"Push item on the heap, then pop and return the smallest item\n\
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)
{
@ -468,6 +515,8 @@ Equivalent to: sorted(iterable)[:n]\n");
static PyMethodDef heapq_methods[] = {
{"heappush", (PyCFunction)heappush,
METH_VARARGS, heappush_doc},
{"heappushpop", (PyCFunction)heappushpop,
METH_VARARGS, heappushpop_doc},
{"heappop", (PyCFunction)heappop,
METH_O, heappop_doc},
{"heapreplace", (PyCFunction)heapreplace,