diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst index 26cffa7c643028a..aa095dbbdcf3e5e 100644 --- a/Doc/library/heapq.rst +++ b/Doc/library/heapq.rst @@ -188,6 +188,44 @@ The module also offers three general purpose functions based on heaps. ``key=str.lower``). Equivalent to: ``sorted(iterable, key=key)[:n]``. +.. class:: MinHeap(iterable=None) + + Return a new min-heap, initialized with the items from *iterable* if provided. + This provides an object-oriented interface to the heap functions. + + .. 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 from the heap. + + .. method:: replace(item) + + Pop and return the smallest item from the heap, and also push the new *item*. + + .. method:: peek() + + Return the smallest item without popping it from the heap. Raises :exc:`IndexError` if the heap is empty. + + .. method:: clear() + + Remove all elements from the heap. + + .. note:: + Iterating over the heap and its representation (``__repr__``) yield items in the order of the underlying heap array, which is not necessarily sorted. + +.. class:: MaxHeap(iterable=None) + + Similar to :class:`MinHeap` but implements a max-heap where the largest element is at the root. + + The latter two functions perform best for smaller values of *n*. For larger values, it is more efficient to use the :func:`sorted` function. Also, when ``n==1``, it is more efficient to use the built-in :func:`min` and :func:`max` diff --git a/Lib/heapq.py b/Lib/heapq.py index a3af6dc05bff377..71c4c4a1b94fd79 100644 --- a/Lib/heapq.py +++ b/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,94 @@ def nlargest(n, iterable, key=None): result.sort(reverse=True) return [elem for (k, order, elem) in result] +class MinHeap: + __slots__ = ('_heap',) + + def __init__(self, iterable=None): + if iterable is None: + self._heap = [] + else: + self._heap = list(iterable) + heapify(self._heap) + + def push(self, item): + heappush(self._heap, item) + + def pop(self): + return heappop(self._heap) + + def pushpop(self, item): + return heappushpop(self._heap, item) + + def replace(self, item): + return heapreplace(self._heap, item) + + def peek(self): + return self._heap[0] + + def clear(self): + self._heap.clear() + + def __len__(self): + return len(self._heap) + + def __bool__(self): + return bool(self._heap) + + def __contains__(self, item): + return item in self._heap + + def __iter__(self): + return iter(self._heap) + + def __repr__(self): + return f'{self.__class__.__name__}({self._heap!r})' + + +class MaxHeap: + __slots__ = ('_heap',) + + def __init__(self, iterable=None): + if iterable is None: + self._heap = [] + else: + self._heap = list(iterable) + heapify_max(self._heap) + + def push(self, item): + heappush_max(self._heap, item) + + def pop(self): + return heappop_max(self._heap) + + def pushpop(self, item): + return heappushpop_max(self._heap, item) + + def replace(self, item): + return heapreplace_max(self._heap, item) + + def peek(self): + return self._heap[0] + + def clear(self): + self._heap.clear() + + def __len__(self): + return len(self._heap) + + def __bool__(self): + return bool(self._heap) + + def __contains__(self, item): + return item in self._heap + + def __iter__(self): + return iter(self._heap) + + def __repr__(self): + return f'{self.__class__.__name__}({self._heap!r})' + + # If available, use C implementation try: from _heapq import * diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index d6623fee9bb2b43..6c5df9fa1431ba6 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -637,5 +637,80 @@ class TestErrorHandlingC(TestErrorHandling, TestCase): module = c_heapq +class TestMinHeap: + def test_init(self): + h = self.module.MinHeap() + self.assertEqual(len(h), 0) + self.assertFalse(h) + + h = self.module.MinHeap([3, 1, 2]) + self.assertEqual(len(h), 3) + self.assertTrue(h) + self.assertEqual(h.peek(), 1) + + def test_push_pop(self): + h = self.module.MinHeap() + h.push(3) + h.push(1) + h.push(2) + self.assertEqual(h.pop(), 1) + self.assertEqual(h.pop(), 2) + self.assertEqual(h.pop(), 3) + with self.assertRaises(IndexError): + h.pop() + + def test_pushpop_replace(self): + h = self.module.MinHeap([2, 3]) + self.assertEqual(h.pushpop(1), 1) + self.assertEqual(len(h), 2) + + self.assertEqual(h.replace(4), 2) + self.assertEqual(h.peek(), 3) + + def test_peek(self): + h = self.module.MinHeap([1, 2, 3]) + self.assertEqual(h.peek(), 1) + h.pop() + self.assertEqual(h.peek(), 2) + h.clear() + with self.assertRaises(IndexError): + h.peek() + + def test_clear_contains_iter(self): + h = self.module.MinHeap([3, 1, 2]) + self.assertIn(1, h) + self.assertNotIn(4, h) + + items = list(h) + self.assertEqual(len(items), 3) + + h.clear() + self.assertEqual(len(h), 0) + self.assertFalse(h) + +class TestMinHeapPython(TestMinHeap, TestCase): + module = py_heapq + +@skipUnless(c_heapq, 'requires _heapq') +class TestMinHeapC(TestMinHeap, TestCase): + module = c_heapq + +class TestMaxHeap: + def test_push_pop(self): + h = self.module.MaxHeap() + h.push(1) + h.push(3) + h.push(2) + self.assertEqual(h.pop(), 3) + self.assertEqual(h.pop(), 2) + self.assertEqual(h.pop(), 1) + +class TestMaxHeapPython(TestMaxHeap, TestCase): + module = py_heapq + +@skipUnless(c_heapq, 'requires _heapq') +class TestMaxHeapC(TestMaxHeap, TestCase): + module = c_heapq + if __name__ == "__main__": unittest.main() diff --git a/Misc/NEWS.d/next/Library/2026-07-08-13-50-00.gh-issue-153366.rst b/Misc/NEWS.d/next/Library/2026-07-08-13-50-00.gh-issue-153366.rst new file mode 100644 index 000000000000000..30ccda503bbc277 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-08-13-50-00.gh-issue-153366.rst @@ -0,0 +1 @@ +Add :class:`MinHeap` and :class:`MaxHeap` classes to the :mod:`heapq` module, providing an object-oriented interface.