diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst index 26cffa7c643028..3349fc6c8f2a0d 100644 --- a/Doc/library/heapq.rst +++ b/Doc/library/heapq.rst @@ -50,6 +50,83 @@ To create a heap, use a list initialized as ``[]``, or transform an existing lis 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 @@ time:: 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:: diff --git a/Doc/whatsnew/3.16.rst b/Doc/whatsnew/3.16.rst index 1a73a79a58b78b..19ca5d9db1710f 100644 --- a/Doc/whatsnew/3.16.rst +++ b/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 -- diff --git a/Lib/heapq.py b/Lib/heapq.py index a3af6dc05bff37..5cf7598abb117f 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,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 * diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index d6623fee9bb2b4..01e575be279a77 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -5,7 +5,7 @@ import doctest from test.support import import_helper -from unittest import TestCase, skipUnless +from unittest import TestCase, mock, skipUnless from operator import itemgetter py_heapq = import_helper.import_fresh_module('heapq', blocked=['_heapq']) @@ -17,6 +17,8 @@ # Add max-heap variants func_names += [func + '_max' for func in func_names] +class_names = ['MinHeap', 'MaxHeap'] + class TestModules(TestCase): def test_py_functions(self): for fname in func_names: @@ -27,6 +29,15 @@ def test_c_functions(self): for fname in func_names: self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq', fname) + def test_py_classes(self): + # MinHeap and MaxHeap are only implemented in pure Python; there is no + # C accelerator for them. They rely on the module-level functions, + # which are the C versions when _heapq is available. + for cname in class_names: + self.assertEqual(getattr(py_heapq, cname).__module__, 'heapq') + if c_heapq is not None: + self.assertEqual(getattr(c_heapq, cname).__module__, 'heapq') + def load_tests(loader, tests, ignore): # The 'merge' function has examples in its docstring which we should test @@ -637,5 +648,213 @@ class TestErrorHandlingC(TestErrorHandling, TestCase): module = c_heapq +#============================================================================== +# Object-oriented interface: MinHeap and MaxHeap. +# +# The classes add no new heap algorithm; they are thin wrappers that delegate +# to the module-level functions. They are tested in two ways: +# +# * Functional "sanity check" tests, run against both the pure-Python and the +# C classes (like the rest of this file). They are end-to-end, so a +# mis-wired method (say, a MaxHeap that pushed with the min-heap routine) +# changes the observable output and is caught even for the C class. +# +# * Delegation tests that patch the module-level functions and check that each +# method forwards to the right one. These only apply to the pure-Python +# class, since the C class calls the underlying routines directly and cannot +# be intercepted; its wiring is covered by the functional tests instead. + + +class HeapClassSanityTest: + """Functional sanity checks shared by the MinHeap/MaxHeap test cases. + + Subclasses set: + + * *cls* -- the heap class under test (e.g. ``py_heapq.MinHeap``) + * *ordered* -- a small list already satisfying the heap invariant + * *reverse* -- ``True`` for max-heaps (the order ``pop()`` yields items in) + """ + + def in_heap_order(self, data): + return sorted(data, reverse=self.reverse) + + def test_construct_empty(self): + h = self.cls() + self.assertEqual(len(h), 0) + self.assertFalse(h) + + def test_construct_from_iterable(self): + data = [random.randrange(1000) for _ in range(100)] + # Accept any iterable, not just a sequence. + h = self.cls(iter(data)) + self.assertEqual(len(h), len(data)) + drained = [h.pop() for _ in range(len(data))] + self.assertEqual(drained, self.in_heap_order(data)) + + def test_push_then_pop_sorts(self): + data = [random.randrange(1000) for _ in range(100)] + h = self.cls() + for item in data: + self.assertIsNone(h.push(item)) + self.assertEqual(len(h), len(data)) + drained = [h.pop() for _ in range(len(data))] + self.assertEqual(drained, self.in_heap_order(data)) + + def test_len_and_bool(self): + h = self.cls() + self.assertEqual(len(h), 0) + self.assertFalse(h) + h.push(1) + self.assertEqual(len(h), 1) + self.assertTrue(h) + + def test_construct_non_iterable_raises(self): + self.assertRaises(TypeError, self.cls, 10) + + def test_pop_empty_raises(self): + h = self.cls() + self.assertRaises(IndexError, h.pop) + + def test_pushpop(self): + h = self.cls(self.ordered) + result = h.pushpop(2) + expected = self.in_heap_order(list(self.ordered) + [2])[0] + self.assertEqual(result, expected) + self.assertEqual(len(h), len(self.ordered)) + + def test_replace(self): + h = self.cls(self.ordered) + expected = self.in_heap_order(self.ordered)[0] + new_item = -1000 if self.reverse else 1000 + self.assertEqual(h.replace(new_item), expected) + self.assertEqual(len(h), len(self.ordered)) + + def test_nsmallest(self): + data = [random.randrange(1000) for _ in range(100)] + h = self.cls(data) + self.assertEqual(h.nsmallest(10), sorted(data)[:10]) + # the heap is not modified + self.assertEqual(len(h), len(data)) + + def test_nlargest(self): + data = [random.randrange(1000) for _ in range(100)] + h = self.cls(data) + self.assertEqual(h.nlargest(10), sorted(data, reverse=True)[:10]) + # the heap is not modified + self.assertEqual(len(h), len(data)) + + def test_nsmallest_nlargest_with_key(self): + data = ['a', 'ee', 'ddd', 'bbbb', 'ccccc'] + h = self.cls(data) + self.assertEqual(h.nsmallest(2, key=len), ['a', 'ee']) + self.assertEqual(h.nlargest(2, key=len), ['ccccc', 'bbbb']) + + def test_repr(self): + h = self.cls(self.ordered) + self.assertEqual(repr(h), f'{self.cls.__name__}({self.ordered!r})') + + +class TestMinHeapPython(HeapClassSanityTest, TestCase): + cls = py_heapq.MinHeap + ordered = [1, 2, 3] + reverse = False + + +@skipUnless(c_heapq, 'requires _heapq') +class TestMinHeapC(HeapClassSanityTest, TestCase): + cls = c_heapq.MinHeap if c_heapq else None + ordered = [1, 2, 3] + reverse = False + + +class TestMaxHeapPython(HeapClassSanityTest, TestCase): + cls = py_heapq.MaxHeap + ordered = [3, 2, 1] + reverse = True + + +@skipUnless(c_heapq, 'requires _heapq') +class TestMaxHeapC(HeapClassSanityTest, TestCase): + cls = c_heapq.MaxHeap if c_heapq else None + ordered = [3, 2, 1] + reverse = True + + +class HeapClassDelegationTest: + """Check that each method forwards to the right module-level function. + + Delegation only makes sense for the pure-Python class, so these tests + always patch ``py_heapq``. + + Subclasses set *cls* and the names of the functions each method is + expected to forward to (*heapify_func*, *push_func*, *pop_func*, + *pushpop_func* and *replace_func*). + """ + + def test_init_delegates_to_heapify(self): + with mock.patch.object(py_heapq, self.heapify_func) as m: + h = self.cls([3, 1, 2]) + m.assert_called_once_with(h._queue) + + def test_init_empty_does_not_heapify(self): + with mock.patch.object(py_heapq, self.heapify_func) as m: + self.cls() + m.assert_not_called() + + def test_push_delegates(self): + h = self.cls() + with mock.patch.object(py_heapq, self.push_func) as m: + h.push(42) + m.assert_called_once_with(h._queue, 42) + + def test_pop_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, self.pop_func) as m: + h.pop() + m.assert_called_once_with(h._queue) + + def test_pushpop_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, self.pushpop_func) as m: + h.pushpop(42) + m.assert_called_once_with(h._queue, 42) + + def test_replace_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, self.replace_func) as m: + h.replace(42) + m.assert_called_once_with(h._queue, 42) + + def test_nsmallest_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, 'nsmallest') as m: + h.nsmallest(2, key=str) + m.assert_called_once_with(2, h._queue, key=str) + + def test_nlargest_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, 'nlargest') as m: + h.nlargest(2, key=str) + m.assert_called_once_with(2, h._queue, key=str) + + +class TestMinHeapDelegation(HeapClassDelegationTest, TestCase): + cls = py_heapq.MinHeap + heapify_func = 'heapify' + push_func = 'heappush' + pop_func = 'heappop' + pushpop_func = 'heappushpop' + replace_func = 'heapreplace' + + +class TestMaxHeapDelegation(HeapClassDelegationTest, TestCase): + cls = py_heapq.MaxHeap + heapify_func = 'heapify_max' + push_func = 'heappush_max' + pop_func = 'heappop_max' + pushpop_func = 'heappushpop_max' + replace_func = 'heapreplace_max' + + if __name__ == "__main__": unittest.main()