Skip to content

Commit f5b7847

Browse files
StanFromIrelandpicnixzencukou
authored
gh-110067: Make max heap methods public and add missing ones (GH-130725)
Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com> Co-authored-by: Petr Viktorin <encukou@gmail.com>
1 parent bb5ec6e commit f5b7847

File tree

7 files changed

+501
-100
lines changed

7 files changed

+501
-100
lines changed

Doc/library/heapq.rst

+88-22
Original file line numberDiff line numberDiff line change
@@ -16,40 +16,56 @@
1616
This module provides an implementation of the heap queue algorithm, also known
1717
as the priority queue algorithm.
1818

19-
Heaps are binary trees for which every parent node has a value less than or
20-
equal to any of its children. We refer to this condition as the heap invariant.
19+
Min-heaps are binary trees for which every parent node has a value less than
20+
or equal to any of its children.
21+
We refer to this condition as the heap invariant.
2122

22-
This implementation uses arrays for which
23-
``heap[k] <= heap[2*k+1]`` and ``heap[k] <= heap[2*k+2]`` for all *k*, counting
24-
elements from zero. For the sake of comparison, non-existing elements are
25-
considered to be infinite. The interesting property of a heap is that its
26-
smallest element is always the root, ``heap[0]``.
23+
For min-heaps, this implementation uses lists for which
24+
``heap[k] <= heap[2*k+1]`` and ``heap[k] <= heap[2*k+2]`` for all *k* for which
25+
the compared elements exist. Elements are counted from zero. The interesting
26+
property of a min-heap is that its smallest element is always the root,
27+
``heap[0]``.
2728

28-
The API below differs from textbook heap algorithms in two aspects: (a) We use
29-
zero-based indexing. This makes the relationship between the index for a node
30-
and the indexes for its children slightly less obvious, but is more suitable
31-
since Python uses zero-based indexing. (b) Our pop method returns the smallest
32-
item, not the largest (called a "min heap" in textbooks; a "max heap" is more
33-
common in texts because of its suitability for in-place sorting).
29+
Max-heaps satisfy the reverse invariant: every parent node has a value
30+
*greater* than any of its children. These are implemented as lists for which
31+
``maxheap[2*k+1] <= maxheap[k]`` and ``maxheap[2*k+2] <= maxheap[k]`` for all
32+
*k* for which the compared elements exist.
33+
The root, ``maxheap[0]``, contains the *largest* element;
34+
``heap.sort(reverse=True)`` maintains the max-heap invariant.
3435

35-
These two make it possible to view the heap as a regular Python list without
36-
surprises: ``heap[0]`` is the smallest item, and ``heap.sort()`` maintains the
37-
heap invariant!
36+
The :mod:`!heapq` API differs from textbook heap algorithms in two aspects: (a)
37+
We use zero-based indexing. This makes the relationship between the index for
38+
a node and the indexes for its children slightly less obvious, but is more
39+
suitable since Python uses zero-based indexing. (b) Textbooks often focus on
40+
max-heaps, due to their suitability for in-place sorting. Our implementation
41+
favors min-heaps as they better correspond to Python :class:`lists <list>`.
3842

39-
To create a heap, use a list initialized to ``[]``, or you can transform a
40-
populated list into a heap via function :func:`heapify`.
43+
These two aspects make it possible to view the heap as a regular Python list
44+
without surprises: ``heap[0]`` is the smallest item, and ``heap.sort()``
45+
maintains the heap invariant!
4146

42-
The following functions are provided:
47+
Like :meth:`list.sort`, this implementation uses only the ``<`` operator
48+
for comparisons, for both min-heaps and max-heaps.
49+
50+
In the API below, and in this documentation, the unqualified term *heap*
51+
generally refers to a min-heap.
52+
The API for max-heaps is named using a ``_max`` suffix.
53+
54+
To create a heap, use a list initialized as ``[]``, or transform an existing list
55+
into a min-heap or max-heap using the :func:`heapify` or :func:`heapify_max`
56+
functions, respectively.
57+
58+
The following functions are provided for min-heaps:
4359

4460

4561
.. function:: heappush(heap, item)
4662

47-
Push the value *item* onto the *heap*, maintaining the heap invariant.
63+
Push the value *item* onto the *heap*, maintaining the min-heap invariant.
4864

4965

5066
.. function:: heappop(heap)
5167

52-
Pop and return the smallest item from the *heap*, maintaining the heap
68+
Pop and return the smallest item from the *heap*, maintaining the min-heap
5369
invariant. If the heap is empty, :exc:`IndexError` is raised. To access the
5470
smallest item without popping it, use ``heap[0]``.
5571

@@ -63,7 +79,7 @@ The following functions are provided:
6379

6480
.. function:: heapify(x)
6581

66-
Transform list *x* into a heap, in-place, in linear time.
82+
Transform list *x* into a min-heap, in-place, in linear time.
6783

6884

6985
.. function:: heapreplace(heap, item)
@@ -82,6 +98,56 @@ The following functions are provided:
8298
on the heap.
8399

84100

101+
For max-heaps, the following functions are provided:
102+
103+
104+
.. function:: heapify_max(x)
105+
106+
Transform list *x* into a max-heap, in-place, in linear time.
107+
108+
.. versionadded:: next
109+
110+
111+
.. function:: heappush_max(heap, item)
112+
113+
Push the value *item* onto the max-heap *heap*, maintaining the max-heap
114+
invariant.
115+
116+
.. versionadded:: next
117+
118+
119+
.. function:: heappop_max(heap)
120+
121+
Pop and return the largest item from the max-heap *heap*, maintaining the
122+
max-heap invariant. If the max-heap is empty, :exc:`IndexError` is raised.
123+
To access the largest item without popping it, use ``maxheap[0]``.
124+
125+
.. versionadded:: next
126+
127+
128+
.. function:: heappushpop_max(heap, item)
129+
130+
Push *item* on the max-heap *heap*, then pop and return the largest item
131+
from *heap*.
132+
The combined action runs more efficiently than :func:`heappush_max`
133+
followed by a separate call to :func:`heappop_max`.
134+
135+
.. versionadded:: next
136+
137+
138+
.. function:: heapreplace_max(heap, item)
139+
140+
Pop and return the largest item from the max-heap *heap* and also push the
141+
new *item*.
142+
The max-heap size doesn't change. If the max-heap is empty,
143+
:exc:`IndexError` is raised.
144+
145+
The value returned may be smaller than the *item* added. Refer to the
146+
analogous function :func:`heapreplace` for detailed usage notes.
147+
148+
.. versionadded:: next
149+
150+
85151
The module also offers three general purpose functions based on heaps.
86152

87153

Doc/whatsnew/3.14.rst

+12
Original file line numberDiff line numberDiff line change
@@ -1114,6 +1114,18 @@ graphlib
11141114
(Contributed by Daniel Pope in :gh:`130914`)
11151115

11161116

1117+
heapq
1118+
-----
1119+
1120+
* Add functions for working with max-heaps:
1121+
1122+
* :func:`heapq.heapify_max`,
1123+
* :func:`heapq.heappush_max`,
1124+
* :func:`heapq.heappop_max`,
1125+
* :func:`heapq.heapreplace_max`
1126+
* :func:`heapq.heappushpop_max`
1127+
1128+
11171129
hmac
11181130
----
11191131

Lib/heapq.py

+29-22
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,7 @@ def heapify(x):
178178
for i in reversed(range(n//2)):
179179
_siftup(x, i)
180180

181-
def _heappop_max(heap):
181+
def heappop_max(heap):
182182
"""Maxheap version of a heappop."""
183183
lastelt = heap.pop() # raises appropriate IndexError if heap is empty
184184
if heap:
@@ -188,19 +188,32 @@ def _heappop_max(heap):
188188
return returnitem
189189
return lastelt
190190

191-
def _heapreplace_max(heap, item):
191+
def heapreplace_max(heap, item):
192192
"""Maxheap version of a heappop followed by a heappush."""
193193
returnitem = heap[0] # raises appropriate IndexError if heap is empty
194194
heap[0] = item
195195
_siftup_max(heap, 0)
196196
return returnitem
197197

198-
def _heapify_max(x):
198+
def heappush_max(heap, item):
199+
"""Maxheap version of a heappush."""
200+
heap.append(item)
201+
_siftdown_max(heap, 0, len(heap)-1)
202+
203+
def heappushpop_max(heap, item):
204+
"""Maxheap fast version of a heappush followed by a heappop."""
205+
if heap and item < heap[0]:
206+
item, heap[0] = heap[0], item
207+
_siftup_max(heap, 0)
208+
return item
209+
210+
def heapify_max(x):
199211
"""Transform list into a maxheap, in-place, in O(len(x)) time."""
200212
n = len(x)
201213
for i in reversed(range(n//2)):
202214
_siftup_max(x, i)
203215

216+
204217
# 'heap' is a heap at all indices >= startpos, except possibly for pos. pos
205218
# is the index of a leaf with a possibly out-of-order value. Restore the
206219
# heap invariant.
@@ -335,9 +348,9 @@ def merge(*iterables, key=None, reverse=False):
335348
h_append = h.append
336349

337350
if reverse:
338-
_heapify = _heapify_max
339-
_heappop = _heappop_max
340-
_heapreplace = _heapreplace_max
351+
_heapify = heapify_max
352+
_heappop = heappop_max
353+
_heapreplace = heapreplace_max
341354
direction = -1
342355
else:
343356
_heapify = heapify
@@ -490,10 +503,10 @@ def nsmallest(n, iterable, key=None):
490503
result = [(elem, i) for i, elem in zip(range(n), it)]
491504
if not result:
492505
return result
493-
_heapify_max(result)
506+
heapify_max(result)
494507
top = result[0][0]
495508
order = n
496-
_heapreplace = _heapreplace_max
509+
_heapreplace = heapreplace_max
497510
for elem in it:
498511
if elem < top:
499512
_heapreplace(result, (elem, order))
@@ -507,10 +520,10 @@ def nsmallest(n, iterable, key=None):
507520
result = [(key(elem), i, elem) for i, elem in zip(range(n), it)]
508521
if not result:
509522
return result
510-
_heapify_max(result)
523+
heapify_max(result)
511524
top = result[0][0]
512525
order = n
513-
_heapreplace = _heapreplace_max
526+
_heapreplace = heapreplace_max
514527
for elem in it:
515528
k = key(elem)
516529
if k < top:
@@ -583,19 +596,13 @@ def nlargest(n, iterable, key=None):
583596
from _heapq import *
584597
except ImportError:
585598
pass
586-
try:
587-
from _heapq import _heapreplace_max
588-
except ImportError:
589-
pass
590-
try:
591-
from _heapq import _heapify_max
592-
except ImportError:
593-
pass
594-
try:
595-
from _heapq import _heappop_max
596-
except ImportError:
597-
pass
598599

600+
# For backwards compatibility
601+
_heappop_max = heappop_max
602+
_heapreplace_max = heapreplace_max
603+
_heappush_max = heappush_max
604+
_heappushpop_max = heappushpop_max
605+
_heapify_max = heapify_max
599606

600607
if __name__ == "__main__":
601608

0 commit comments

Comments
 (0)