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

@ -45,6 +45,13 @@ The following functions are provided:
Pop and return the smallest item from the *heap*, maintaining the heap Pop and return the smallest item from the *heap*, maintaining the heap
invariant. If the heap is empty, :exc:`IndexError` is raised. invariant. If the heap is empty, :exc:`IndexError` is raised.
.. function:: heappushpop(heap, item)
Push *item* on the heap, then pop and return the smallest item from the
*heap*. The combined action runs more efficiently than :func:`heappush`
followed by a separate call to :func:`heappop`.
.. versionadded:: 2.6
.. function:: heapify(x) .. function:: heapify(x)

View File

@ -127,7 +127,7 @@ From all times, sorting has always been a Great Art! :-)
""" """
__all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'merge', __all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'merge',
'nlargest', 'nsmallest'] 'nlargest', 'nsmallest', 'heappushpop']
from itertools import islice, repeat, count, imap, izip, tee from itertools import islice, repeat, count, imap, izip, tee
from operator import itemgetter, neg from operator import itemgetter, neg
@ -165,6 +165,13 @@ def heapreplace(heap, item):
_siftup(heap, 0) _siftup(heap, 0)
return returnitem return returnitem
def heappushpop(heap, item):
"""Fast version of a heappush followed by a heappop."""
if heap and item > heap[0]:
item, heap[0] = heap[0], item
_siftup(heap, 0)
return item
def heapify(x): def heapify(x):
"""Transform list into a heap, in-place, in O(len(heap)) time.""" """Transform list into a heap, in-place, in O(len(heap)) time."""
n = len(x) n = len(x)
@ -304,7 +311,7 @@ def _siftup(heap, pos):
# If available, use C implementation # If available, use C implementation
try: try:
from _heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest from _heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest, heappushpop
except ImportError: except ImportError:
pass pass

View File

@ -107,6 +107,34 @@ class TestHeap(unittest.TestCase):
self.assertRaises(TypeError, self.module.heapreplace, None, None) self.assertRaises(TypeError, self.module.heapreplace, None, None)
self.assertRaises(IndexError, self.module.heapreplace, [], None) self.assertRaises(IndexError, self.module.heapreplace, [], None)
def test_nbest_with_pushpop(self):
data = [random.randrange(2000) for i in range(1000)]
heap = data[:10]
self.module.heapify(heap)
for item in data[10:]:
self.module.heappushpop(heap, item)
self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:])
self.assertEqual(self.module.heappushpop([], 'x'), 'x')
def test_heappushpop(self):
h = []
x = self.module.heappushpop(h, 10)
self.assertEqual((h, x), ([], 10))
h = [10]
x = self.module.heappushpop(h, 10.0)
self.assertEqual((h, x), ([10], 10.0))
self.assertEqual(type(h[0]), int)
self.assertEqual(type(x), float)
h = [10];
x = self.module.heappushpop(h, 9)
self.assertEqual((h, x), ([10], 9))
h = [10];
x = self.module.heappushpop(h, 11)
self.assertEqual((h, x), ([11], 10))
def test_heapsort(self): def test_heapsort(self):
# Exercise everything with repeated heapsort checks # Exercise everything with repeated heapsort checks
for trial in xrange(100): for trial in xrange(100):

View File

@ -491,6 +491,8 @@ Core and builtins
Library Library
------- -------
- #2274 Add heapq.heappushpop().
- Add inspect.isabstract(object) to fix bug #2223 - Add inspect.isabstract(object) to fix bug #2223
- Add a __format__ method to Decimal, to support PEP 3101. - Add a __format__ method to Decimal, to support PEP 3101.

View File

@ -162,6 +162,11 @@ heapreplace(PyObject *self, PyObject *args)
{ {
PyObject *heap, *item, *returnitem; 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)) if (!PyArg_UnpackTuple(args, "heapreplace", 2, 2, &heap, &item))
return NULL; return NULL;
@ -195,6 +200,48 @@ this routine unless written as part of a conditional replacement:\n\n\
if item > heap[0]:\n\ if item > heap[0]:\n\
item = heapreplace(heap, item)\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 * static PyObject *
heapify(PyObject *self, PyObject *heap) heapify(PyObject *self, PyObject *heap)
{ {
@ -468,6 +515,8 @@ Equivalent to: sorted(iterable)[:n]\n");
static PyMethodDef heapq_methods[] = { static PyMethodDef heapq_methods[] = {
{"heappush", (PyCFunction)heappush, {"heappush", (PyCFunction)heappush,
METH_VARARGS, heappush_doc}, METH_VARARGS, heappush_doc},
{"heappushpop", (PyCFunction)heappushpop,
METH_VARARGS, heappushpop_doc},
{"heappop", (PyCFunction)heappop, {"heappop", (PyCFunction)heappop,
METH_O, heappop_doc}, METH_O, heappop_doc},
{"heapreplace", (PyCFunction)heapreplace, {"heapreplace", (PyCFunction)heapreplace,