gh-110067: Make max heap methods public and add missing ones

Doc/library/heapq.rst

@@ -82,6 +82,55 @@ The following functions are provided:
    on the heap.
 
 
+.. function:: heapify_max(x)
+
+   Transform list *x* into a heap, in-place, in linear time.
+
+   .. versionadded:: next
+
+
+.. function:: heappush_max(heap, item)
+
+   Push the value *item* onto the *heap*, maintaining the heap invariant.
+
+   .. versionadded:: next
+
+.. function:: heappop_max(heap)
+
+   Pop and return the largest item from the *heap*, maintaining the heap
+   invariant.  If the heap is empty, :exc:`IndexError` is raised.  To access the
+   largest item without popping it, use ``heap[0]``.
+
+   .. versionadded:: next
+
+
+.. function:: heappushpop_max(heap, item)
+
+   Push *item* on the heap, then pop and return the largest item from the
+   *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 *heap*, and also push the new *item*.
+   The heap size doesn't change. If the heap is empty, :exc:`IndexError` is raised.
+
+   This one step operation is more efficient than a :func:`heappop` followed by
+   :func:`heappush` and can be more appropriate when using a fixed-size heap.
+   The pop/push combination always returns an element from the heap and replaces
+   it with *item*.
+
+   The value returned may be larger than the *item* added.  If that isn't
+   desired, consider using :func:`heappushpop` instead.  Its push/pop
+   combination returns the smaller of the two values, leaving the larger value
+   on the heap.
+
+   .. versionadded:: next
+
+
 The module also offers three general purpose functions based on heaps.
 
 

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,44 @@ 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 heap[0] < item:
+        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)
 
+# For backwards compadibility
+def _heappop_max(heap):
+    return heappop_max(heap)
+def _heapreplace_max(heap, item):
+    return heapreplace_max(heap, item)
+def _heappush_max(heap, item):
+    return _heappush_max(heap, item)
+def _heappushpop_max(heap, item):
+    return _heappushpop_max(heap, item)
+def _heapify_max(x):
+    return _heapify_max()
+
+
 # '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 +360,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 +515,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 +532,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,18 +608,6 @@ def nlargest(n, iterable, key=None):
     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
 
 
 if __name__ == "__main__":

Lib/test/test_heapq.py

@@ -14,7 +14,7 @@
 # _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']
+              'heappop_max', 'heapreplace_max', 'heapify_max', 'heappushpop_max',]
 
 class TestModules(TestCase):
     def test_py_functions(self):
@@ -74,13 +74,47 @@ def test_push_pop(self):
         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, [])
+        try:
+            self.assertRaises(TypeError, self.module.heappush, None, None)
+            self.assertRaises(TypeError, self.module.heappop, None)
+        except AttributeError:
+            pass
+
     def check_invariant(self, heap):
         # Check the heap invariant.
         for pos, item in enumerate(heap):
             if pos: # pos 0 has no parent
                 parentpos = (pos-1) >> 1
                 self.assertTrue(heap[parentpos] <= item)
 
+    def check_max_invariant(self, heap):
+        for pos in range(1, len(heap)):
+            parentpos = (pos - 1) >> 1
+            self.assertTrue(heap[parentpos] >= heap[pos])
+
     def test_heapify(self):
         for size in list(range(30)) + [20000]:
             heap = [random.random() for dummy in range(size)]
@@ -89,6 +123,14 @@ def test_heapify(self):
 
         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, None)
+
     def test_naive_nbest(self):
         data = [random.randrange(2000) for i in range(1000)]
         heap = []
@@ -153,12 +195,31 @@ def test_heappushpop(self):
         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.assertEqual((h, x), ([], 10))
+
+        h = [10]
+        x = self.module.heappushpop_max(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_max(h, 11)
+        self.assertEqual((h, x), ([11], 10))
+
+        h = [10]
+        x = self.module.heappushpop_max(h, 9)
+        self.assertEqual((h, x), ([10], 9))
+
     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

Misc/NEWS.d/next/Library/2025-03-01-15-00-00.gh-issue-110067.1ad3as.rst

@@ -0,0 +1 @@
+Make max heap functions public.

Modules/_heapqmodule.c

@@ -482,7 +482,7 @@ siftup_max(PyListObject *heap, Py_ssize_t pos)
 
 
 /*[clinic input]
-_heapq._heappop_max
+_heapq.heappop_max
 
     heap: object(subclass_of='&PyList_Type')
     /
@@ -491,14 +491,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 +508,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 +524,63 @@ 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(top, item, Py_LT);
+    Py_DECREF(top);
+    if (cmp < 0)
+        return NULL;
+    if (cmp == 0) {
+        return Py_NewRef(item);
+    }
+    returnitem = PyList_GET_ITEM(heap, 0);
+    PyList_SET_ITEM(heap, 0, Py_NewRef(item));
+
+    if (siftup_max((PyListObject *)heap, 0)) {
+        Py_DECREF(returnitem);
+        return NULL;
+    }
+    return Py_NewRef(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_HEAPPOP_MAX_METHODDEF
+    _HEAPQ_HEAPIFY_MAX_METHODDEF
+    _HEAPQ_HEAPREPLACE_MAX_METHODDEF
+    _HEAPQ_HEAPPUSHPOP_MAX_METHODDEF
     {NULL, NULL}           /* sentinel */
 };