From 53bdf093437349907807da9143f9c2bdcea9ab3a Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Thu, 13 Mar 2008 19:03:51 +0000 Subject: [PATCH] Issue 2274: Add heapq.heappushpop(). --- Doc/library/heapq.rst | 7 ++++++ Lib/heapq.py | 11 ++++++++-- Lib/test/test_heapq.py | 28 ++++++++++++++++++++++++ Misc/NEWS | 2 ++ Modules/_heapqmodule.c | 49 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 95 insertions(+), 2 deletions(-) diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst index 115d2236921..1168fb688a2 100644 --- a/Doc/library/heapq.rst +++ b/Doc/library/heapq.rst @@ -45,6 +45,13 @@ The following functions are provided: Pop and return the smallest item from the *heap*, maintaining the heap 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) diff --git a/Lib/heapq.py b/Lib/heapq.py index 39e38008822..23f5fcbfda0 100644 --- a/Lib/heapq.py +++ b/Lib/heapq.py @@ -127,7 +127,7 @@ From all times, sorting has always been a Great Art! :-) """ __all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'merge', - 'nlargest', 'nsmallest'] + 'nlargest', 'nsmallest', 'heappushpop'] from itertools import islice, repeat, count, imap, izip, tee from operator import itemgetter, neg @@ -165,6 +165,13 @@ def heapreplace(heap, item): _siftup(heap, 0) 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): """Transform list into a heap, in-place, in O(len(heap)) time.""" n = len(x) @@ -304,7 +311,7 @@ def _siftup(heap, pos): # If available, use C implementation try: - from _heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest + from _heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest, heappushpop except ImportError: pass diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index 7f6e4b59e70..fec027e3b65 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -107,6 +107,34 @@ class TestHeap(unittest.TestCase): self.assertRaises(TypeError, self.module.heapreplace, None, 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): # Exercise everything with repeated heapsort checks for trial in xrange(100): diff --git a/Misc/NEWS b/Misc/NEWS index bd3a6cca0dc..4d0b9f0fb52 100644 --- a/Misc/NEWS +++ b/Misc/NEWS @@ -491,6 +491,8 @@ Core and builtins Library ------- +- #2274 Add heapq.heappushpop(). + - Add inspect.isabstract(object) to fix bug #2223 - Add a __format__ method to Decimal, to support PEP 3101. diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c index f970cae9137..703742e2608 100644 --- a/Modules/_heapqmodule.c +++ b/Modules/_heapqmodule.c @@ -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,