gh-153357: Add object-oriented MinHeap and MaxHeap classes to heapq

Doc/library/heapq.rst

@@ -50,6 +50,83 @@
 into a min-heap or max-heap using the :func:`heapify` or :func:`heapify_max`
 functions, respectively.
 
+Alternatively, the :class:`MinHeap` and :class:`MaxHeap` classes offer an
+object-oriented interface, described next; the module-level functions that they
+build upon are documented afterwards.
+
+
+Heap classes
+------------
+
+The :class:`MinHeap` and :class:`MaxHeap` classes wrap the module's heap
+operations in an object that owns its own storage.  Compared with the functions
+below, there is no separate list to pass to every call or to accidentally
+break the heap invariant with an unrelated list operation.
+
+.. class:: MinHeap(iterable=None)
+           MaxHeap(iterable=None)
+
+   Create a new heap.  If *iterable* is given, the heap is initialized from its
+   items in linear time; otherwise it starts empty.
+
+   A :class:`MinHeap` keeps its smallest item at the top, while a
+   :class:`MaxHeap` keeps its largest item at the top.  The two classes share
+   the same interface and differ only in that ordering --- that is, in which
+   item :meth:`~MinHeap.pop` removes and which item the other methods treat as
+   the top.  As with the functions below, only the ``<`` operator is used to
+   compare items.
+
+   The methods documented below belong to both classes; a :class:`MinHeap` is
+   used in the descriptions, with *MaxHeap* behaving the same way once
+   "smallest" is read as "largest".
+
+   .. method:: push(item)
+
+      Push the value *item* onto the heap, maintaining the heap invariant.
+
+   .. method:: pop()
+
+      Pop and return the smallest item from the heap, maintaining the heap
+      invariant.  If the heap is empty, :exc:`IndexError` is raised.
+
+   .. method:: pushpop(item)
+
+      Push *item* on the heap, then pop and return the smallest item.  The
+      combined action runs more efficiently than :meth:`push` followed by a
+      separate call to :meth:`pop`.
+
+   .. method:: replace(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 :meth:`pop` followed by
+      :meth:`push` and can be more appropriate when using a fixed-size heap.
+      The value returned may be larger than the *item* added; consider using
+      :meth:`pushpop` instead if that is not desired.
+
+   .. method:: nsmallest(n, key=None)
+
+      Return a list with the *n* smallest items from the heap, from smallest to
+      largest, without modifying the heap.  Behaves like the module-level
+      :func:`nsmallest` function applied to the heap's items.
+
+   .. method:: nlargest(n, key=None)
+
+      Return a list with the *n* largest items from the heap, from largest to
+      smallest, without modifying the heap.  Behaves like the module-level
+      :func:`nlargest` function applied to the heap's items.
+
+   In addition, ``len(heap)`` returns the number of items on the heap, and an
+   empty heap is falsy, so ``bool(heap)`` is ``False`` when it holds no items.
+
+   .. versionadded:: 3.16
+
+
+Heap functions
+--------------
+
 The following functions are provided for min-heaps:
 
 
@@ -214,6 +291,15 @@
 This is similar to ``sorted(iterable)``, but unlike :func:`sorted`, this
 implementation is not stable.
 
+Using the object-oriented interface, the same sort becomes::
+
+   >>> def heapsort(iterable):
+   ...     heap = MinHeap(iterable)
+   ...     return [heap.pop() for _ in range(len(heap))]
+   ...
+   >>> heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
+   [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+
 Heap elements can be tuples.  This is useful for assigning comparison values
 (such as task priorities) alongside the main record being tracked::
 

Doc/whatsnew/3.16.rst

@@ -221,6 +221,13 @@ gzip
   which is passed on to the constructor of the :class:`~gzip.GzipFile` class.
   (Contributed by Marin Misur in :gh:`91372`.)
 
+heapq
+-----
+
+* Add the :class:`~heapq.MinHeap` and :class:`~heapq.MaxHeap` classes, an
+  object-oriented interface to the module's heap functions.
+  (Contributed by Facundo Batista in :gh:`137200`.)
+
 io
 --
 

Lib/heapq.py

@@ -128,7 +128,8 @@
 
 __all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'heappushpop',
            'heappush_max', 'heappop_max', 'heapify_max', 'heapreplace_max',
-           'heappushpop_max', 'nlargest', 'nsmallest', 'merge']
+           'heappushpop_max', 'nlargest', 'nsmallest', 'merge',
+           'MinHeap', 'MaxHeap']
 
 def heappush(heap, item):
     """Push item onto heap, maintaining the heap invariant."""
@@ -592,6 +593,159 @@ def nlargest(n, iterable, key=None):
     result.sort(reverse=True)
     return [elem for (k, order, elem) in result]
 
+
+class MinHeap:
+    """Min-heap: a priority queue that serves its smallest item first.
+
+    Only the ``<`` operator is used to compare items.
+
+    If *iterable* is given, the heap is initialized from its items in linear
+    time; otherwise it starts empty.
+
+    >>> h = MinHeap([3, 1, 2])
+    >>> h.push(0)
+    >>> [h.pop() for _ in range(len(h))]
+    [0, 1, 2, 3]
+    """
+
+    def __init__(self, iterable=None):
+        if iterable is None:
+            self._queue = []
+        else:
+            self._queue = list(iterable)
+            heapify(self._queue)
+
+    def push(self, item):
+        """Push item onto the heap, maintaining the heap invariant."""
+        heappush(self._queue, item)
+
+    def pop(self):
+        """Pop and return the smallest item, maintaining the heap invariant.
+
+        Raises IndexError if the heap is empty.
+        """
+        return heappop(self._queue)
+
+    def pushpop(self, item):
+        """Push item on the heap, then pop and return the smallest item.
+
+        The combined action runs more efficiently than a push() followed by
+        a separate call to pop().
+        """
+        return heappushpop(self._queue, item)
+
+    def replace(self, item):
+        """Pop and return the smallest item, and also push the new item.
+
+        The heap size doesn't change.  Raises IndexError if the heap is
+        empty.  This is more efficient than a pop() followed by a push(),
+        and can be more appropriate when using a fixed-size heap.  Note that
+        the value returned may be larger than item.
+        """
+        return heapreplace(self._queue, item)
+
+    def nsmallest(self, n, key=None):
+        """Return a list of the n smallest items, smallest first.
+
+        The heap is left unchanged.  If *key* is given, it is applied to
+        each item to determine the ordering.
+        """
+        return nsmallest(n, self._queue, key=key)
+
+    def nlargest(self, n, key=None):
+        """Return a list of the n largest items, largest first.
+
+        The heap is left unchanged.  If *key* is given, it is applied to
+        each item to determine the ordering.
+        """
+        return nlargest(n, self._queue, key=key)
+
+    def __len__(self):
+        return len(self._queue)
+
+    def __bool__(self):
+        return bool(self._queue)
+
+    def __repr__(self):
+        return f'{type(self).__name__}({self._queue!r})'
+
+
+class MaxHeap:
+    """Max-heap: a priority queue that serves its largest item first.
+
+    Only the ``<`` operator is used to compare items.
+
+    If *iterable* is given, the heap is initialized from its items in linear
+    time; otherwise it starts empty.
+
+    >>> h = MaxHeap([1, 3, 2])
+    >>> h.push(4)
+    >>> [h.pop() for _ in range(len(h))]
+    [4, 3, 2, 1]
+    """
+
+    def __init__(self, iterable=None):
+        if iterable is None:
+            self._queue = []
+        else:
+            self._queue = list(iterable)
+            heapify_max(self._queue)
+
+    def push(self, item):
+        """Push item onto the heap, maintaining the heap invariant."""
+        heappush_max(self._queue, item)
+
+    def pop(self):
+        """Pop and return the largest item, maintaining the heap invariant.
+
+        Raises IndexError if the heap is empty.
+        """
+        return heappop_max(self._queue)
+
+    def pushpop(self, item):
+        """Push item on the heap, then pop and return the largest item.
+
+        The combined action runs more efficiently than a push() followed by
+        a separate call to pop().
+        """
+        return heappushpop_max(self._queue, item)
+
+    def replace(self, item):
+        """Pop and return the largest item, and also push the new item.
+
+        The heap size doesn't change.  Raises IndexError if the heap is
+        empty.  This is more efficient than a pop() followed by a push(),
+        and can be more appropriate when using a fixed-size heap.  Note that
+        the value returned may be smaller than item.
+        """
+        return heapreplace_max(self._queue, item)
+
+    def nsmallest(self, n, key=None):
+        """Return a list of the n smallest items, smallest first.
+
+        The heap is left unchanged.  If *key* is given, it is applied to
+        each item to determine the ordering.
+        """
+        return nsmallest(n, self._queue, key=key)
+
+    def nlargest(self, n, key=None):
+        """Return a list of the n largest items, largest first.
+
+        The heap is left unchanged.  If *key* is given, it is applied to
+        each item to determine the ordering.
+        """
+        return nlargest(n, self._queue, key=key)
+
+    def __len__(self):
+        return len(self._queue)
+
+    def __bool__(self):
+        return bool(self._queue)
+
+    def __repr__(self):
+        return f'{type(self).__name__}({self._queue!r})'
+
+
 # If available, use C implementation
 try:
     from _heapq import *