Initial addition

Author: StanFromIreland

Doc/library/heapq.rst

@@ -82,6 +82,46 @@ The following functions are provided:
    on the heap.
 
 
+.. function:: heappush_max(heap, item)
+
+   Push the value *item* onto the *heap*, maintaining the heap invariant.
+
+
+.. 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]``.
+
+
+.. 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`.
+
+
+.. function:: heapify_max(x)
+
+   Transform list *x* into a max heap, in-place, in linear time.
+
+
+.. function:: heapreplace(heap, item)
+
+   Pop and return the smallest 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.
+
+
 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,14 +188,26 @@ 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)):
@@ -335,9 +347,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 +502,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 +519,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:
@@ -584,15 +596,15 @@ def nlargest(n, iterable, key=None):
 except ImportError:
     pass
 try:
-    from _heapq import _heapreplace_max
+    from _heapq import heapreplace_max
 except ImportError:
     pass
 try:
-    from _heapq import _heapify_max
+    from _heapq import heapify_max
 except ImportError:
     pass
 try:
-    from _heapq import _heappop_max
+    from _heapq import heappop_max
 except ImportError:
     pass
 

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']
 
 class TestModules(TestCase):
     def test_py_functions(self):
@@ -23,7 +23,7 @@ def test_py_functions(self):
 
     @skipUnless(c_heapq, 'requires _heapq')
     def test_c_functions(self):
-        for fname in func_names:
+        for fname in ['heapify', 'heappop', 'heappush', 'heappushpop', 'heapreplace']:
             self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq')
 
 
@@ -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.