diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst index d3c4b920ba5..2bd0162a982 100644 --- a/Doc/library/heapq.rst +++ b/Doc/library/heapq.rst @@ -16,40 +16,56 @@ This module provides an implementation of the heap queue algorithm, also known as the priority queue algorithm. -Heaps are binary trees for which every parent node has a value less than or -equal to any of its children. We refer to this condition as the heap invariant. +Min-heaps are binary trees for which every parent node has a value less than +or equal to any of its children. +We refer to this condition as the heap invariant. -This implementation uses arrays for which -``heap[k] <= heap[2*k+1]`` and ``heap[k] <= heap[2*k+2]`` for all *k*, counting -elements from zero. For the sake of comparison, non-existing elements are -considered to be infinite. The interesting property of a heap is that its -smallest element is always the root, ``heap[0]``. +For min-heaps, this implementation uses lists for which +``heap[k] <= heap[2*k+1]`` and ``heap[k] <= heap[2*k+2]`` for all *k* for which +the compared elements exist. Elements are counted from zero. The interesting +property of a min-heap is that its smallest element is always the root, +``heap[0]``. -The API below differs from textbook heap algorithms in two aspects: (a) We use -zero-based indexing. This makes the relationship between the index for a node -and the indexes for its children slightly less obvious, but is more suitable -since Python uses zero-based indexing. (b) Our pop method returns the smallest -item, not the largest (called a "min heap" in textbooks; a "max heap" is more -common in texts because of its suitability for in-place sorting). +Max-heaps satisfy the reverse invariant: every parent node has a value +*greater* than any of its children. These are implemented as lists for which +``maxheap[2*k+1] <= maxheap[k]`` and ``maxheap[2*k+2] <= maxheap[k]`` for all +*k* for which the compared elements exist. +The root, ``maxheap[0]``, contains the *largest* element; +``heap.sort(reverse=True)`` maintains the max-heap invariant. -These two make it possible to view the heap as a regular Python list without -surprises: ``heap[0]`` is the smallest item, and ``heap.sort()`` maintains the -heap invariant! +The :mod:`!heapq` API differs from textbook heap algorithms in two aspects: (a) +We use zero-based indexing. This makes the relationship between the index for +a node and the indexes for its children slightly less obvious, but is more +suitable since Python uses zero-based indexing. (b) Textbooks often focus on +max-heaps, due to their suitability for in-place sorting. Our implementation +favors min-heaps as they better correspond to Python :class:`lists `. -To create a heap, use a list initialized to ``[]``, or you can transform a -populated list into a heap via function :func:`heapify`. +These two aspects make it possible to view the heap as a regular Python list +without surprises: ``heap[0]`` is the smallest item, and ``heap.sort()`` +maintains the heap invariant! -The following functions are provided: +Like :meth:`list.sort`, this implementation uses only the ``<`` operator +for comparisons, for both min-heaps and max-heaps. + +In the API below, and in this documentation, the unqualified term *heap* +generally refers to a min-heap. +The API for max-heaps is named using a ``_max`` suffix. + +To create a heap, use a list initialized as ``[]``, or transform an existing list +into a min-heap or max-heap using the :func:`heapify` or :func:`heapify_max` +functions, respectively. + +The following functions are provided for min-heaps: .. function:: heappush(heap, item) - Push the value *item* onto the *heap*, maintaining the heap invariant. + Push the value *item* onto the *heap*, maintaining the min-heap invariant. .. function:: heappop(heap) - Pop and return the smallest item from the *heap*, maintaining the heap + Pop and return the smallest item from the *heap*, maintaining the min-heap invariant. If the heap is empty, :exc:`IndexError` is raised. To access the smallest item without popping it, use ``heap[0]``. @@ -63,7 +79,7 @@ The following functions are provided: .. function:: heapify(x) - Transform list *x* into a heap, in-place, in linear time. + Transform list *x* into a min-heap, in-place, in linear time. .. function:: heapreplace(heap, item) @@ -82,6 +98,56 @@ The following functions are provided: on the heap. +For max-heaps, the following functions are provided: + + +.. function:: heapify_max(x) + + Transform list *x* into a max-heap, in-place, in linear time. + + .. versionadded:: next + + +.. function:: heappush_max(heap, item) + + Push the value *item* onto the max-heap *heap*, maintaining the max-heap + invariant. + + .. versionadded:: next + + +.. function:: heappop_max(heap) + + Pop and return the largest item from the max-heap *heap*, maintaining the + max-heap invariant. If the max-heap is empty, :exc:`IndexError` is raised. + To access the largest item without popping it, use ``maxheap[0]``. + + .. versionadded:: next + + +.. function:: heappushpop_max(heap, item) + + Push *item* on the max-heap *heap*, then pop and return the largest item + from *heap*. + The combined action runs more efficiently than :func:`heappush_max` + followed by a separate call to :func:`heappop_max`. + + .. versionadded:: next + + +.. function:: heapreplace_max(heap, item) + + Pop and return the largest item from the max-heap *heap* and also push the + new *item*. + The max-heap size doesn't change. If the max-heap is empty, + :exc:`IndexError` is raised. + + The value returned may be smaller than the *item* added. Refer to the + analogous function :func:`heapreplace` for detailed usage notes. + + .. versionadded:: next + + The module also offers three general purpose functions based on heaps. diff --git a/Doc/whatsnew/3.14.rst b/Doc/whatsnew/3.14.rst index f09e517d456..506a820175b 100644 --- a/Doc/whatsnew/3.14.rst +++ b/Doc/whatsnew/3.14.rst @@ -1114,6 +1114,18 @@ graphlib (Contributed by Daniel Pope in :gh:`130914`) +heapq +----- + +* Add functions for working with max-heaps: + + * :func:`heapq.heapify_max`, + * :func:`heapq.heappush_max`, + * :func:`heapq.heappop_max`, + * :func:`heapq.heapreplace_max` + * :func:`heapq.heappushpop_max` + + hmac ---- diff --git a/Lib/heapq.py b/Lib/heapq.py index 9649da251f2..6ceb211f1ca 100644 --- a/Lib/heapq.py +++ b/Lib/heapq.py @@ -178,7 +178,7 @@ def heapify(x): for i in reversed(range(n//2)): _siftup(x, i) -def _heappop_max(heap): +def heappop_max(heap): """Maxheap version of a heappop.""" lastelt = heap.pop() # raises appropriate IndexError if heap is empty if heap: @@ -188,19 +188,32 @@ def _heappop_max(heap): return returnitem return lastelt -def _heapreplace_max(heap, item): +def heapreplace_max(heap, item): """Maxheap version of a heappop followed by a heappush.""" returnitem = heap[0] # raises appropriate IndexError if heap is empty heap[0] = item _siftup_max(heap, 0) return returnitem -def _heapify_max(x): +def heappush_max(heap, item): + """Maxheap version of a heappush.""" + heap.append(item) + _siftdown_max(heap, 0, len(heap)-1) + +def heappushpop_max(heap, item): + """Maxheap fast version of a heappush followed by a heappop.""" + if heap and item < heap[0]: + item, heap[0] = heap[0], item + _siftup_max(heap, 0) + return item + +def heapify_max(x): """Transform list into a maxheap, in-place, in O(len(x)) time.""" n = len(x) for i in reversed(range(n//2)): _siftup_max(x, i) + # 'heap' is a heap at all indices >= startpos, except possibly for pos. pos # is the index of a leaf with a possibly out-of-order value. Restore the # heap invariant. @@ -335,9 +348,9 @@ def merge(*iterables, key=None, reverse=False): h_append = h.append if reverse: - _heapify = _heapify_max - _heappop = _heappop_max - _heapreplace = _heapreplace_max + _heapify = heapify_max + _heappop = heappop_max + _heapreplace = heapreplace_max direction = -1 else: _heapify = heapify @@ -490,10 +503,10 @@ def nsmallest(n, iterable, key=None): result = [(elem, i) for i, elem in zip(range(n), it)] if not result: return result - _heapify_max(result) + heapify_max(result) top = result[0][0] order = n - _heapreplace = _heapreplace_max + _heapreplace = heapreplace_max for elem in it: if elem < top: _heapreplace(result, (elem, order)) @@ -507,10 +520,10 @@ def nsmallest(n, iterable, key=None): result = [(key(elem), i, elem) for i, elem in zip(range(n), it)] if not result: return result - _heapify_max(result) + heapify_max(result) top = result[0][0] order = n - _heapreplace = _heapreplace_max + _heapreplace = heapreplace_max for elem in it: k = key(elem) if k < top: @@ -583,19 +596,13 @@ try: from _heapq import * except ImportError: pass -try: - from _heapq import _heapreplace_max -except ImportError: - pass -try: - from _heapq import _heapify_max -except ImportError: - pass -try: - from _heapq import _heappop_max -except ImportError: - pass +# For backwards compatibility +_heappop_max = heappop_max +_heapreplace_max = heapreplace_max +_heappush_max = heappush_max +_heappushpop_max = heappushpop_max +_heapify_max = heapify_max if __name__ == "__main__": diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index 1aa8e4e2897..d6623fee9bb 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -13,8 +13,9 @@ c_heapq = import_helper.import_fresh_module('heapq', fresh=['_heapq']) # _heapq.nlargest/nsmallest are saved in heapq._nlargest/_smallest when # _heapq is imported, so check them there -func_names = ['heapify', 'heappop', 'heappush', 'heappushpop', 'heapreplace', - '_heappop_max', '_heapreplace_max', '_heapify_max'] +func_names = ['heapify', 'heappop', 'heappush', 'heappushpop', 'heapreplace'] +# Add max-heap variants +func_names += [func + '_max' for func in func_names] class TestModules(TestCase): def test_py_functions(self): @@ -24,7 +25,7 @@ class TestModules(TestCase): @skipUnless(c_heapq, 'requires _heapq') def test_c_functions(self): for fname in func_names: - self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq') + self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq', fname) def load_tests(loader, tests, ignore): @@ -74,6 +75,34 @@ class TestHeap: except AttributeError: pass + def test_max_push_pop(self): + # 1) Push 256 random numbers and pop them off, verifying all's OK. + heap = [] + data = [] + self.check_max_invariant(heap) + for i in range(256): + item = random.random() + data.append(item) + self.module.heappush_max(heap, item) + self.check_max_invariant(heap) + results = [] + while heap: + item = self.module.heappop_max(heap) + self.check_max_invariant(heap) + results.append(item) + data_sorted = data[:] + data_sorted.sort(reverse=True) + + self.assertEqual(data_sorted, results) + # 2) Check that the invariant holds for a sorted array + self.check_max_invariant(results) + + self.assertRaises(TypeError, self.module.heappush_max, []) + + exc_types = (AttributeError, TypeError) + self.assertRaises(exc_types, self.module.heappush_max, None, None) + self.assertRaises(exc_types, self.module.heappop_max, None) + def check_invariant(self, heap): # Check the heap invariant. for pos, item in enumerate(heap): @@ -81,6 +110,11 @@ class TestHeap: parentpos = (pos-1) >> 1 self.assertTrue(heap[parentpos] <= item) + def check_max_invariant(self, heap): + for pos, item in enumerate(heap[1:], start=1): + parentpos = (pos - 1) >> 1 + self.assertGreaterEqual(heap[parentpos], item) + def test_heapify(self): for size in list(range(30)) + [20000]: heap = [random.random() for dummy in range(size)] @@ -89,6 +123,14 @@ class TestHeap: self.assertRaises(TypeError, self.module.heapify, None) + def test_heapify_max(self): + for size in list(range(30)) + [20000]: + heap = [random.random() for dummy in range(size)] + self.module.heapify_max(heap) + self.check_max_invariant(heap) + + self.assertRaises(TypeError, self.module.heapify_max, None) + def test_naive_nbest(self): data = [random.randrange(2000) for i in range(1000)] heap = [] @@ -109,10 +151,7 @@ class TestHeap: def test_nbest(self): # Less-naive "N-best" algorithm, much faster (if len(data) is big - # enough ) than sorting all of data. However, if we had a max - # heap instead of a min heap, it could go faster still via - # heapify'ing all of data (linear time), then doing 10 heappops - # (10 log-time steps). + # enough ) than sorting all of data. data = [random.randrange(2000) for i in range(1000)] heap = data[:10] self.module.heapify(heap) @@ -125,6 +164,17 @@ class TestHeap: self.assertRaises(TypeError, self.module.heapreplace, None, None) self.assertRaises(IndexError, self.module.heapreplace, [], None) + def test_nbest_maxheap(self): + # With a max heap instead of a min heap, the "N-best" algorithm can + # go even faster still via heapify'ing all of data (linear time), then + # doing 10 heappops (10 log-time steps). + data = [random.randrange(2000) for i in range(1000)] + heap = data[:] + self.module.heapify_max(heap) + result = [self.module.heappop_max(heap) for _ in range(10)] + result.reverse() + self.assertEqual(result, sorted(data)[-10:]) + def test_nbest_with_pushpop(self): data = [random.randrange(2000) for i in range(1000)] heap = data[:10] @@ -134,6 +184,62 @@ class TestHeap: self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:]) self.assertEqual(self.module.heappushpop([], 'x'), 'x') + def test_naive_nworst(self): + # Max-heap variant of "test_naive_nbest" + data = [random.randrange(2000) for i in range(1000)] + heap = [] + for item in data: + self.module.heappush_max(heap, item) + if len(heap) > 10: + self.module.heappop_max(heap) + heap.sort() + expected = sorted(data)[:10] + self.assertEqual(heap, expected) + + def heapiter_max(self, heap): + # An iterator returning a max-heap's elements, largest-first. + try: + while 1: + yield self.module.heappop_max(heap) + except IndexError: + pass + + def test_nworst(self): + # Max-heap variant of "test_nbest" + data = [random.randrange(2000) for i in range(1000)] + heap = data[:10] + self.module.heapify_max(heap) + for item in data[10:]: + if item < heap[0]: # this gets rarer the longer we run + self.module.heapreplace_max(heap, item) + expected = sorted(data, reverse=True)[-10:] + self.assertEqual(list(self.heapiter_max(heap)), expected) + + self.assertRaises(TypeError, self.module.heapreplace_max, None) + self.assertRaises(TypeError, self.module.heapreplace_max, None, None) + self.assertRaises(IndexError, self.module.heapreplace_max, [], None) + + def test_nworst_minheap(self): + # Min-heap variant of "test_nbest_maxheap" + data = [random.randrange(2000) for i in range(1000)] + heap = data[:] + self.module.heapify(heap) + result = [self.module.heappop(heap) for _ in range(10)] + result.reverse() + expected = sorted(data, reverse=True)[-10:] + self.assertEqual(result, expected) + + def test_nworst_with_pushpop(self): + # Max-heap variant of "test_nbest_with_pushpop" + data = [random.randrange(2000) for i in range(1000)] + heap = data[:10] + self.module.heapify_max(heap) + for item in data[10:]: + self.module.heappushpop_max(heap, item) + expected = sorted(data, reverse=True)[-10:] + self.assertEqual(list(self.heapiter_max(heap)), expected) + self.assertEqual(self.module.heappushpop_max([], 'x'), 'x') + def test_heappushpop(self): h = [] x = self.module.heappushpop(h, 10) @@ -153,12 +259,31 @@ class TestHeap: x = self.module.heappushpop(h, 11) self.assertEqual((h, x), ([11], 10)) + def test_heappushpop_max(self): + h = [] + x = self.module.heappushpop_max(h, 10) + self.assertTupleEqual((h, x), ([], 10)) + + h = [10] + x = self.module.heappushpop_max(h, 10.0) + self.assertTupleEqual((h, x), ([10], 10.0)) + self.assertIsInstance(h[0], int) + self.assertIsInstance(x, float) + + h = [10] + x = self.module.heappushpop_max(h, 11) + self.assertTupleEqual((h, x), ([10], 11)) + + h = [10] + x = self.module.heappushpop_max(h, 9) + self.assertTupleEqual((h, x), ([9], 10)) + def test_heappop_max(self): - # _heapop_max has an optimization for one-item lists which isn't + # heapop_max has an optimization for one-item lists which isn't # covered in other tests, so test that case explicitly here h = [3, 2] - self.assertEqual(self.module._heappop_max(h), 3) - self.assertEqual(self.module._heappop_max(h), 2) + self.assertEqual(self.module.heappop_max(h), 3) + self.assertEqual(self.module.heappop_max(h), 2) def test_heapsort(self): # Exercise everything with repeated heapsort checks @@ -175,6 +300,20 @@ class TestHeap: heap_sorted = [self.module.heappop(heap) for i in range(size)] self.assertEqual(heap_sorted, sorted(data)) + def test_heapsort_max(self): + for trial in range(100): + size = random.randrange(50) + data = [random.randrange(25) for i in range(size)] + if trial & 1: # Half of the time, use heapify_max + heap = data[:] + self.module.heapify_max(heap) + else: # The rest of the time, use heappush_max + heap = [] + for item in data: + self.module.heappush_max(heap, item) + heap_sorted = [self.module.heappop_max(heap) for i in range(size)] + self.assertEqual(heap_sorted, sorted(data, reverse=True)) + def test_merge(self): inputs = [] for i in range(random.randrange(25)): @@ -377,16 +516,20 @@ class SideEffectLT: class TestErrorHandling: def test_non_sequence(self): - for f in (self.module.heapify, self.module.heappop): + for f in (self.module.heapify, self.module.heappop, + self.module.heapify_max, self.module.heappop_max): self.assertRaises((TypeError, AttributeError), f, 10) for f in (self.module.heappush, self.module.heapreplace, + self.module.heappush_max, self.module.heapreplace_max, self.module.nlargest, self.module.nsmallest): self.assertRaises((TypeError, AttributeError), f, 10, 10) def test_len_only(self): - for f in (self.module.heapify, self.module.heappop): + for f in (self.module.heapify, self.module.heappop, + self.module.heapify_max, self.module.heappop_max): self.assertRaises((TypeError, AttributeError), f, LenOnly()) - for f in (self.module.heappush, self.module.heapreplace): + for f in (self.module.heappush, self.module.heapreplace, + self.module.heappush_max, self.module.heapreplace_max): self.assertRaises((TypeError, AttributeError), f, LenOnly(), 10) for f in (self.module.nlargest, self.module.nsmallest): self.assertRaises(TypeError, f, 2, LenOnly()) @@ -395,7 +538,8 @@ class TestErrorHandling: seq = [CmpErr(), CmpErr(), CmpErr()] for f in (self.module.heapify, self.module.heappop): self.assertRaises(ZeroDivisionError, f, seq) - for f in (self.module.heappush, self.module.heapreplace): + for f in (self.module.heappush, self.module.heapreplace, + self.module.heappush_max, self.module.heapreplace_max): self.assertRaises(ZeroDivisionError, f, seq, 10) for f in (self.module.nlargest, self.module.nsmallest): self.assertRaises(ZeroDivisionError, f, 2, seq) @@ -403,6 +547,8 @@ class TestErrorHandling: def test_arg_parsing(self): for f in (self.module.heapify, self.module.heappop, self.module.heappush, self.module.heapreplace, + self.module.heapify_max, self.module.heappop_max, + self.module.heappush_max, self.module.heapreplace_max, self.module.nlargest, self.module.nsmallest): self.assertRaises((TypeError, AttributeError), f, 10) @@ -424,6 +570,10 @@ class TestErrorHandling: # Python version raises IndexError, C version RuntimeError with self.assertRaises((IndexError, RuntimeError)): self.module.heappush(heap, SideEffectLT(5, heap)) + heap = [] + heap.extend(SideEffectLT(i, heap) for i in range(200)) + with self.assertRaises((IndexError, RuntimeError)): + self.module.heappush_max(heap, SideEffectLT(5, heap)) def test_heappop_mutating_heap(self): heap = [] @@ -431,8 +581,12 @@ class TestErrorHandling: # Python version raises IndexError, C version RuntimeError with self.assertRaises((IndexError, RuntimeError)): self.module.heappop(heap) + heap = [] + heap.extend(SideEffectLT(i, heap) for i in range(200)) + with self.assertRaises((IndexError, RuntimeError)): + self.module.heappop_max(heap) - def test_comparison_operator_modifiying_heap(self): + def test_comparison_operator_modifying_heap(self): # See bpo-39421: Strong references need to be taken # when comparing objects as they can alter the heap class EvilClass(int): @@ -444,7 +598,7 @@ class TestErrorHandling: self.module.heappush(heap, EvilClass(0)) self.assertRaises(IndexError, self.module.heappushpop, heap, 1) - def test_comparison_operator_modifiying_heap_two_heaps(self): + def test_comparison_operator_modifying_heap_two_heaps(self): class h(int): def __lt__(self, o): @@ -464,6 +618,17 @@ class TestErrorHandling: self.assertRaises((IndexError, RuntimeError), self.module.heappush, list1, g(1)) self.assertRaises((IndexError, RuntimeError), self.module.heappush, list2, h(1)) + list1, list2 = [], [] + + self.module.heappush_max(list1, h(0)) + self.module.heappush_max(list2, g(0)) + self.module.heappush_max(list1, g(1)) + self.module.heappush_max(list2, h(1)) + + self.assertRaises((IndexError, RuntimeError), self.module.heappush_max, list1, g(1)) + self.assertRaises((IndexError, RuntimeError), self.module.heappush_max, list2, h(1)) + + class TestErrorHandlingPython(TestErrorHandling, TestCase): module = py_heapq diff --git a/Misc/NEWS.d/next/Library/2025-03-01-15-00-00.gh-issue-110067.1ad3as.rst b/Misc/NEWS.d/next/Library/2025-03-01-15-00-00.gh-issue-110067.1ad3as.rst new file mode 100644 index 00000000000..98e125f4966 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2025-03-01-15-00-00.gh-issue-110067.1ad3as.rst @@ -0,0 +1,5 @@ +Make :mod:`heapq` max-heap functions :func:`heapq.heapify_max`, :func:`heapq.heappush_max`, +:func:`heapq.heappop_max`, and :func:`heapq.heapreplace_max` public. +Previous underscored naming is kept for backwards compatibility. +Additionally, the missing function :func:`heapq.heappushpop_max` has been added +to both the C and Python implementations. diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c index 80fe9cff985..095866eec7d 100644 --- a/Modules/_heapqmodule.c +++ b/Modules/_heapqmodule.c @@ -480,9 +480,33 @@ siftup_max(PyListObject *heap, Py_ssize_t pos) return siftdown_max(heap, startpos, pos); } +/*[clinic input] +_heapq.heappush_max + + heap: object(subclass_of='&PyList_Type') + item: object + / + +Push item onto max heap, maintaining the heap invariant. +[clinic start generated code]*/ + +static PyObject * +_heapq_heappush_max_impl(PyObject *module, PyObject *heap, PyObject *item) +/*[clinic end generated code: output=c869d5f9deb08277 input=4743d7db137b6e2b]*/ +{ + if (PyList_Append(heap, item)) { + return NULL; + } + + if (siftdown_max((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-1)) { + return NULL; + } + + Py_RETURN_NONE; +} /*[clinic input] -_heapq._heappop_max +_heapq.heappop_max heap: object(subclass_of='&PyList_Type') / @@ -491,14 +515,14 @@ Maxheap variant of heappop. [clinic start generated code]*/ static PyObject * -_heapq__heappop_max_impl(PyObject *module, PyObject *heap) -/*[clinic end generated code: output=9e77aadd4e6a8760 input=362c06e1c7484793]*/ +_heapq_heappop_max_impl(PyObject *module, PyObject *heap) +/*[clinic end generated code: output=2f051195ab404b77 input=e62b14016a5a26de]*/ { return heappop_internal(heap, siftup_max); } /*[clinic input] -_heapq._heapreplace_max +_heapq.heapreplace_max heap: object(subclass_of='&PyList_Type') item: object @@ -508,15 +532,14 @@ Maxheap variant of heapreplace. [clinic start generated code]*/ static PyObject * -_heapq__heapreplace_max_impl(PyObject *module, PyObject *heap, - PyObject *item) -/*[clinic end generated code: output=8ad7545e4a5e8adb input=f2dd27cbadb948d7]*/ +_heapq_heapreplace_max_impl(PyObject *module, PyObject *heap, PyObject *item) +/*[clinic end generated code: output=8770778b5a9cbe9b input=21a3d28d757c881c]*/ { return heapreplace_internal(heap, item, siftup_max); } /*[clinic input] -_heapq._heapify_max +_heapq.heapify_max heap: object(subclass_of='&PyList_Type') / @@ -525,21 +548,74 @@ Maxheap variant of heapify. [clinic start generated code]*/ static PyObject * -_heapq__heapify_max_impl(PyObject *module, PyObject *heap) -/*[clinic end generated code: output=2cb028beb4a8b65e input=c1f765ee69f124b8]*/ +_heapq_heapify_max_impl(PyObject *module, PyObject *heap) +/*[clinic end generated code: output=8401af3856529807 input=edda4255728c431e]*/ { return heapify_internal(heap, siftup_max); } +/*[clinic input] +_heapq.heappushpop_max + + heap: object(subclass_of='&PyList_Type') + item: object + / + +Maxheap variant of heappushpop. + +The combined action runs more efficiently than heappush_max() followed by +a separate call to heappop_max(). +[clinic start generated code]*/ + +static PyObject * +_heapq_heappushpop_max_impl(PyObject *module, PyObject *heap, PyObject *item) +/*[clinic end generated code: output=ff0019f0941aca0d input=525a843013cbd6c0]*/ +{ + PyObject *returnitem; + int cmp; + + if (PyList_GET_SIZE(heap) == 0) { + return Py_NewRef(item); + } + + PyObject *top = PyList_GET_ITEM(heap, 0); + Py_INCREF(top); + cmp = PyObject_RichCompareBool(item, top, Py_LT); + Py_DECREF(top); + if (cmp < 0) { + return NULL; + } + if (cmp == 0) { + return Py_NewRef(item); + } + + if (PyList_GET_SIZE(heap) == 0) { + PyErr_SetString(PyExc_IndexError, "index out of range"); + return NULL; + } + + returnitem = PyList_GET_ITEM(heap, 0); + PyList_SET_ITEM(heap, 0, Py_NewRef(item)); + if (siftup_max((PyListObject *)heap, 0) < 0) { + Py_DECREF(returnitem); + return NULL; + } + return returnitem; +} + static PyMethodDef heapq_methods[] = { _HEAPQ_HEAPPUSH_METHODDEF _HEAPQ_HEAPPUSHPOP_METHODDEF _HEAPQ_HEAPPOP_METHODDEF _HEAPQ_HEAPREPLACE_METHODDEF _HEAPQ_HEAPIFY_METHODDEF - _HEAPQ__HEAPPOP_MAX_METHODDEF - _HEAPQ__HEAPIFY_MAX_METHODDEF - _HEAPQ__HEAPREPLACE_MAX_METHODDEF + + _HEAPQ_HEAPPUSH_MAX_METHODDEF + _HEAPQ_HEAPPUSHPOP_MAX_METHODDEF + _HEAPQ_HEAPPOP_MAX_METHODDEF + _HEAPQ_HEAPREPLACE_MAX_METHODDEF + _HEAPQ_HEAPIFY_MAX_METHODDEF + {NULL, NULL} /* sentinel */ }; diff --git a/Modules/clinic/_heapqmodule.c.h b/Modules/clinic/_heapqmodule.c.h index 90463079907..81d10862726 100644 --- a/Modules/clinic/_heapqmodule.c.h +++ b/Modules/clinic/_heapqmodule.c.h @@ -175,96 +175,166 @@ exit: return return_value; } -PyDoc_STRVAR(_heapq__heappop_max__doc__, -"_heappop_max($module, heap, /)\n" +PyDoc_STRVAR(_heapq_heappush_max__doc__, +"heappush_max($module, heap, item, /)\n" "--\n" "\n" -"Maxheap variant of heappop."); +"Push item onto max heap, maintaining the heap invariant."); -#define _HEAPQ__HEAPPOP_MAX_METHODDEF \ - {"_heappop_max", (PyCFunction)_heapq__heappop_max, METH_O, _heapq__heappop_max__doc__}, +#define _HEAPQ_HEAPPUSH_MAX_METHODDEF \ + {"heappush_max", _PyCFunction_CAST(_heapq_heappush_max), METH_FASTCALL, _heapq_heappush_max__doc__}, static PyObject * -_heapq__heappop_max_impl(PyObject *module, PyObject *heap); +_heapq_heappush_max_impl(PyObject *module, PyObject *heap, PyObject *item); static PyObject * -_heapq__heappop_max(PyObject *module, PyObject *arg) -{ - PyObject *return_value = NULL; - PyObject *heap; - - if (!PyList_Check(arg)) { - _PyArg_BadArgument("_heappop_max", "argument", "list", arg); - goto exit; - } - heap = arg; - return_value = _heapq__heappop_max_impl(module, heap); - -exit: - return return_value; -} - -PyDoc_STRVAR(_heapq__heapreplace_max__doc__, -"_heapreplace_max($module, heap, item, /)\n" -"--\n" -"\n" -"Maxheap variant of heapreplace."); - -#define _HEAPQ__HEAPREPLACE_MAX_METHODDEF \ - {"_heapreplace_max", _PyCFunction_CAST(_heapq__heapreplace_max), METH_FASTCALL, _heapq__heapreplace_max__doc__}, - -static PyObject * -_heapq__heapreplace_max_impl(PyObject *module, PyObject *heap, - PyObject *item); - -static PyObject * -_heapq__heapreplace_max(PyObject *module, PyObject *const *args, Py_ssize_t nargs) +_heapq_heappush_max(PyObject *module, PyObject *const *args, Py_ssize_t nargs) { PyObject *return_value = NULL; PyObject *heap; PyObject *item; - if (!_PyArg_CheckPositional("_heapreplace_max", nargs, 2, 2)) { + if (!_PyArg_CheckPositional("heappush_max", nargs, 2, 2)) { goto exit; } if (!PyList_Check(args[0])) { - _PyArg_BadArgument("_heapreplace_max", "argument 1", "list", args[0]); + _PyArg_BadArgument("heappush_max", "argument 1", "list", args[0]); goto exit; } heap = args[0]; item = args[1]; - return_value = _heapq__heapreplace_max_impl(module, heap, item); + return_value = _heapq_heappush_max_impl(module, heap, item); exit: return return_value; } -PyDoc_STRVAR(_heapq__heapify_max__doc__, -"_heapify_max($module, heap, /)\n" +PyDoc_STRVAR(_heapq_heappop_max__doc__, +"heappop_max($module, heap, /)\n" "--\n" "\n" -"Maxheap variant of heapify."); +"Maxheap variant of heappop."); -#define _HEAPQ__HEAPIFY_MAX_METHODDEF \ - {"_heapify_max", (PyCFunction)_heapq__heapify_max, METH_O, _heapq__heapify_max__doc__}, +#define _HEAPQ_HEAPPOP_MAX_METHODDEF \ + {"heappop_max", (PyCFunction)_heapq_heappop_max, METH_O, _heapq_heappop_max__doc__}, static PyObject * -_heapq__heapify_max_impl(PyObject *module, PyObject *heap); +_heapq_heappop_max_impl(PyObject *module, PyObject *heap); static PyObject * -_heapq__heapify_max(PyObject *module, PyObject *arg) +_heapq_heappop_max(PyObject *module, PyObject *arg) { PyObject *return_value = NULL; PyObject *heap; if (!PyList_Check(arg)) { - _PyArg_BadArgument("_heapify_max", "argument", "list", arg); + _PyArg_BadArgument("heappop_max", "argument", "list", arg); goto exit; } heap = arg; - return_value = _heapq__heapify_max_impl(module, heap); + return_value = _heapq_heappop_max_impl(module, heap); exit: return return_value; } -/*[clinic end generated code: output=05f2afdf3bc54c9d input=a9049054013a1b77]*/ + +PyDoc_STRVAR(_heapq_heapreplace_max__doc__, +"heapreplace_max($module, heap, item, /)\n" +"--\n" +"\n" +"Maxheap variant of heapreplace."); + +#define _HEAPQ_HEAPREPLACE_MAX_METHODDEF \ + {"heapreplace_max", _PyCFunction_CAST(_heapq_heapreplace_max), METH_FASTCALL, _heapq_heapreplace_max__doc__}, + +static PyObject * +_heapq_heapreplace_max_impl(PyObject *module, PyObject *heap, PyObject *item); + +static PyObject * +_heapq_heapreplace_max(PyObject *module, PyObject *const *args, Py_ssize_t nargs) +{ + PyObject *return_value = NULL; + PyObject *heap; + PyObject *item; + + if (!_PyArg_CheckPositional("heapreplace_max", nargs, 2, 2)) { + goto exit; + } + if (!PyList_Check(args[0])) { + _PyArg_BadArgument("heapreplace_max", "argument 1", "list", args[0]); + goto exit; + } + heap = args[0]; + item = args[1]; + return_value = _heapq_heapreplace_max_impl(module, heap, item); + +exit: + return return_value; +} + +PyDoc_STRVAR(_heapq_heapify_max__doc__, +"heapify_max($module, heap, /)\n" +"--\n" +"\n" +"Maxheap variant of heapify."); + +#define _HEAPQ_HEAPIFY_MAX_METHODDEF \ + {"heapify_max", (PyCFunction)_heapq_heapify_max, METH_O, _heapq_heapify_max__doc__}, + +static PyObject * +_heapq_heapify_max_impl(PyObject *module, PyObject *heap); + +static PyObject * +_heapq_heapify_max(PyObject *module, PyObject *arg) +{ + PyObject *return_value = NULL; + PyObject *heap; + + if (!PyList_Check(arg)) { + _PyArg_BadArgument("heapify_max", "argument", "list", arg); + goto exit; + } + heap = arg; + return_value = _heapq_heapify_max_impl(module, heap); + +exit: + return return_value; +} + +PyDoc_STRVAR(_heapq_heappushpop_max__doc__, +"heappushpop_max($module, heap, item, /)\n" +"--\n" +"\n" +"Maxheap variant of heappushpop.\n" +"\n" +"The combined action runs more efficiently than heappush_max() followed by\n" +"a separate call to heappop_max()."); + +#define _HEAPQ_HEAPPUSHPOP_MAX_METHODDEF \ + {"heappushpop_max", _PyCFunction_CAST(_heapq_heappushpop_max), METH_FASTCALL, _heapq_heappushpop_max__doc__}, + +static PyObject * +_heapq_heappushpop_max_impl(PyObject *module, PyObject *heap, PyObject *item); + +static PyObject * +_heapq_heappushpop_max(PyObject *module, PyObject *const *args, Py_ssize_t nargs) +{ + PyObject *return_value = NULL; + PyObject *heap; + PyObject *item; + + if (!_PyArg_CheckPositional("heappushpop_max", nargs, 2, 2)) { + goto exit; + } + if (!PyList_Check(args[0])) { + _PyArg_BadArgument("heappushpop_max", "argument 1", "list", args[0]); + goto exit; + } + heap = args[0]; + item = args[1]; + return_value = _heapq_heappushpop_max_impl(module, heap, item); + +exit: + return return_value; +} +/*[clinic end generated code: output=f55d8595ce150c76 input=a9049054013a1b77]*/