mirror of
https://github.com/python/cpython.git
synced 2025-07-07 11:25:30 +00:00
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>
This commit is contained in:
parent
bb5ec6ea6e
commit
f5b784741d
7 changed files with 524 additions and 123 deletions
|
@ -16,40 +16,56 @@
|
|||
This module provides an implementation of the heap queue algorithm, also known
|
||||
as the priority queue algorithm.
|
||||
|
||||
Heaps are binary trees for which every parent node has a value less than or
|
||||
equal to any of its children. We refer to this condition as the heap invariant.
|
||||
Min-heaps are binary trees for which every parent node has a value less than
|
||||
or equal to any of its children.
|
||||
We refer to this condition as the heap invariant.
|
||||
|
||||
This implementation uses arrays for which
|
||||
``heap[k] <= heap[2*k+1]`` and ``heap[k] <= heap[2*k+2]`` for all *k*, counting
|
||||
elements from zero. For the sake of comparison, non-existing elements are
|
||||
considered to be infinite. The interesting property of a heap is that its
|
||||
smallest element is always the root, ``heap[0]``.
|
||||
For min-heaps, this implementation uses lists for which
|
||||
``heap[k] <= heap[2*k+1]`` and ``heap[k] <= heap[2*k+2]`` for all *k* for which
|
||||
the compared elements exist. Elements are counted from zero. The interesting
|
||||
property of a min-heap is that its smallest element is always the root,
|
||||
``heap[0]``.
|
||||
|
||||
The API below differs from textbook heap algorithms in two aspects: (a) We use
|
||||
zero-based indexing. This makes the relationship between the index for a node
|
||||
and the indexes for its children slightly less obvious, but is more suitable
|
||||
since Python uses zero-based indexing. (b) Our pop method returns the smallest
|
||||
item, not the largest (called a "min heap" in textbooks; a "max heap" is more
|
||||
common in texts because of its suitability for in-place sorting).
|
||||
Max-heaps satisfy the reverse invariant: every parent node has a value
|
||||
*greater* than any of its children. These are implemented as lists for which
|
||||
``maxheap[2*k+1] <= maxheap[k]`` and ``maxheap[2*k+2] <= maxheap[k]`` for all
|
||||
*k* for which the compared elements exist.
|
||||
The root, ``maxheap[0]``, contains the *largest* element;
|
||||
``heap.sort(reverse=True)`` maintains the max-heap invariant.
|
||||
|
||||
These two make it possible to view the heap as a regular Python list without
|
||||
surprises: ``heap[0]`` is the smallest item, and ``heap.sort()`` maintains the
|
||||
heap invariant!
|
||||
The :mod:`!heapq` API differs from textbook heap algorithms in two aspects: (a)
|
||||
We use zero-based indexing. This makes the relationship between the index for
|
||||
a node and the indexes for its children slightly less obvious, but is more
|
||||
suitable since Python uses zero-based indexing. (b) Textbooks often focus on
|
||||
max-heaps, due to their suitability for in-place sorting. Our implementation
|
||||
favors min-heaps as they better correspond to Python :class:`lists <list>`.
|
||||
|
||||
To create a heap, use a list initialized to ``[]``, or you can transform a
|
||||
populated list into a heap via function :func:`heapify`.
|
||||
These two aspects make it possible to view the heap as a regular Python list
|
||||
without surprises: ``heap[0]`` is the smallest item, and ``heap.sort()``
|
||||
maintains the heap invariant!
|
||||
|
||||
The following functions are provided:
|
||||
Like :meth:`list.sort`, this implementation uses only the ``<`` operator
|
||||
for comparisons, for both min-heaps and max-heaps.
|
||||
|
||||
In the API below, and in this documentation, the unqualified term *heap*
|
||||
generally refers to a min-heap.
|
||||
The API for max-heaps is named using a ``_max`` suffix.
|
||||
|
||||
To create a heap, use a list initialized as ``[]``, or transform an existing list
|
||||
into a min-heap or max-heap using the :func:`heapify` or :func:`heapify_max`
|
||||
functions, respectively.
|
||||
|
||||
The following functions are provided for min-heaps:
|
||||
|
||||
|
||||
.. function:: heappush(heap, item)
|
||||
|
||||
Push the value *item* onto the *heap*, maintaining the heap invariant.
|
||||
Push the value *item* onto the *heap*, maintaining the min-heap invariant.
|
||||
|
||||
|
||||
.. function:: heappop(heap)
|
||||
|
||||
Pop and return the smallest item from the *heap*, maintaining the heap
|
||||
Pop and return the smallest item from the *heap*, maintaining the min-heap
|
||||
invariant. If the heap is empty, :exc:`IndexError` is raised. To access the
|
||||
smallest item without popping it, use ``heap[0]``.
|
||||
|
||||
|
@ -63,7 +79,7 @@ The following functions are provided:
|
|||
|
||||
.. function:: heapify(x)
|
||||
|
||||
Transform list *x* into a heap, in-place, in linear time.
|
||||
Transform list *x* into a min-heap, in-place, in linear time.
|
||||
|
||||
|
||||
.. function:: heapreplace(heap, item)
|
||||
|
@ -82,6 +98,56 @@ The following functions are provided:
|
|||
on the heap.
|
||||
|
||||
|
||||
For max-heaps, the following functions are provided:
|
||||
|
||||
|
||||
.. function:: heapify_max(x)
|
||||
|
||||
Transform list *x* into a max-heap, in-place, in linear time.
|
||||
|
||||
.. versionadded:: next
|
||||
|
||||
|
||||
.. function:: heappush_max(heap, item)
|
||||
|
||||
Push the value *item* onto the max-heap *heap*, maintaining the max-heap
|
||||
invariant.
|
||||
|
||||
.. versionadded:: next
|
||||
|
||||
|
||||
.. function:: heappop_max(heap)
|
||||
|
||||
Pop and return the largest item from the max-heap *heap*, maintaining the
|
||||
max-heap invariant. If the max-heap is empty, :exc:`IndexError` is raised.
|
||||
To access the largest item without popping it, use ``maxheap[0]``.
|
||||
|
||||
.. versionadded:: next
|
||||
|
||||
|
||||
.. function:: heappushpop_max(heap, item)
|
||||
|
||||
Push *item* on the max-heap *heap*, then pop and return the largest item
|
||||
from *heap*.
|
||||
The combined action runs more efficiently than :func:`heappush_max`
|
||||
followed by a separate call to :func:`heappop_max`.
|
||||
|
||||
.. versionadded:: next
|
||||
|
||||
|
||||
.. function:: heapreplace_max(heap, item)
|
||||
|
||||
Pop and return the largest item from the max-heap *heap* and also push the
|
||||
new *item*.
|
||||
The max-heap size doesn't change. If the max-heap is empty,
|
||||
:exc:`IndexError` is raised.
|
||||
|
||||
The value returned may be smaller than the *item* added. Refer to the
|
||||
analogous function :func:`heapreplace` for detailed usage notes.
|
||||
|
||||
.. versionadded:: next
|
||||
|
||||
|
||||
The module also offers three general purpose functions based on heaps.
|
||||
|
||||
|
||||
|
|
|
@ -1114,6 +1114,18 @@ graphlib
|
|||
(Contributed by Daniel Pope in :gh:`130914`)
|
||||
|
||||
|
||||
heapq
|
||||
-----
|
||||
|
||||
* Add functions for working with max-heaps:
|
||||
|
||||
* :func:`heapq.heapify_max`,
|
||||
* :func:`heapq.heappush_max`,
|
||||
* :func:`heapq.heappop_max`,
|
||||
* :func:`heapq.heapreplace_max`
|
||||
* :func:`heapq.heappushpop_max`
|
||||
|
||||
|
||||
hmac
|
||||
----
|
||||
|
||||
|
|
51
Lib/heapq.py
51
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,19 +188,32 @@ 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 item < heap[0]:
|
||||
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)):
|
||||
_siftup_max(x, i)
|
||||
|
||||
|
||||
# 'heap' is a heap at all indices >= startpos, except possibly for pos. pos
|
||||
# is the index of a leaf with a possibly out-of-order value. Restore the
|
||||
# heap invariant.
|
||||
|
@ -335,9 +348,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 +503,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 +520,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:
|
||||
|
@ -583,19 +596,13 @@ try:
|
|||
from _heapq import *
|
||||
except ImportError:
|
||||
pass
|
||||
try:
|
||||
from _heapq import _heapreplace_max
|
||||
except ImportError:
|
||||
pass
|
||||
try:
|
||||
from _heapq import _heapify_max
|
||||
except ImportError:
|
||||
pass
|
||||
try:
|
||||
from _heapq import _heappop_max
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# For backwards compatibility
|
||||
_heappop_max = heappop_max
|
||||
_heapreplace_max = heapreplace_max
|
||||
_heappush_max = heappush_max
|
||||
_heappushpop_max = heappushpop_max
|
||||
_heapify_max = heapify_max
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
|
|
|
@ -13,8 +13,9 @@ c_heapq = import_helper.import_fresh_module('heapq', fresh=['_heapq'])
|
|||
|
||||
# _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']
|
||||
func_names = ['heapify', 'heappop', 'heappush', 'heappushpop', 'heapreplace']
|
||||
# Add max-heap variants
|
||||
func_names += [func + '_max' for func in func_names]
|
||||
|
||||
class TestModules(TestCase):
|
||||
def test_py_functions(self):
|
||||
|
@ -24,7 +25,7 @@ class TestModules(TestCase):
|
|||
@skipUnless(c_heapq, 'requires _heapq')
|
||||
def test_c_functions(self):
|
||||
for fname in func_names:
|
||||
self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq')
|
||||
self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq', fname)
|
||||
|
||||
|
||||
def load_tests(loader, tests, ignore):
|
||||
|
@ -74,6 +75,34 @@ class TestHeap:
|
|||
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_max, [])
|
||||
|
||||
exc_types = (AttributeError, TypeError)
|
||||
self.assertRaises(exc_types, self.module.heappush_max, None, None)
|
||||
self.assertRaises(exc_types, self.module.heappop_max, None)
|
||||
|
||||
def check_invariant(self, heap):
|
||||
# Check the heap invariant.
|
||||
for pos, item in enumerate(heap):
|
||||
|
@ -81,6 +110,11 @@ class TestHeap:
|
|||
parentpos = (pos-1) >> 1
|
||||
self.assertTrue(heap[parentpos] <= item)
|
||||
|
||||
def check_max_invariant(self, heap):
|
||||
for pos, item in enumerate(heap[1:], start=1):
|
||||
parentpos = (pos - 1) >> 1
|
||||
self.assertGreaterEqual(heap[parentpos], item)
|
||||
|
||||
def test_heapify(self):
|
||||
for size in list(range(30)) + [20000]:
|
||||
heap = [random.random() for dummy in range(size)]
|
||||
|
@ -89,6 +123,14 @@ class TestHeap:
|
|||
|
||||
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_max, None)
|
||||
|
||||
def test_naive_nbest(self):
|
||||
data = [random.randrange(2000) for i in range(1000)]
|
||||
heap = []
|
||||
|
@ -109,10 +151,7 @@ class TestHeap:
|
|||
|
||||
def test_nbest(self):
|
||||
# Less-naive "N-best" algorithm, much faster (if len(data) is big
|
||||
# enough <wink>) than sorting all of data. However, if we had a max
|
||||
# heap instead of a min heap, it could go faster still via
|
||||
# heapify'ing all of data (linear time), then doing 10 heappops
|
||||
# (10 log-time steps).
|
||||
# enough <wink>) than sorting all of data.
|
||||
data = [random.randrange(2000) for i in range(1000)]
|
||||
heap = data[:10]
|
||||
self.module.heapify(heap)
|
||||
|
@ -125,6 +164,17 @@ class TestHeap:
|
|||
self.assertRaises(TypeError, self.module.heapreplace, None, None)
|
||||
self.assertRaises(IndexError, self.module.heapreplace, [], None)
|
||||
|
||||
def test_nbest_maxheap(self):
|
||||
# With a max heap instead of a min heap, the "N-best" algorithm can
|
||||
# go even faster still via heapify'ing all of data (linear time), then
|
||||
# doing 10 heappops (10 log-time steps).
|
||||
data = [random.randrange(2000) for i in range(1000)]
|
||||
heap = data[:]
|
||||
self.module.heapify_max(heap)
|
||||
result = [self.module.heappop_max(heap) for _ in range(10)]
|
||||
result.reverse()
|
||||
self.assertEqual(result, sorted(data)[-10:])
|
||||
|
||||
def test_nbest_with_pushpop(self):
|
||||
data = [random.randrange(2000) for i in range(1000)]
|
||||
heap = data[:10]
|
||||
|
@ -134,6 +184,62 @@ class TestHeap:
|
|||
self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:])
|
||||
self.assertEqual(self.module.heappushpop([], 'x'), 'x')
|
||||
|
||||
def test_naive_nworst(self):
|
||||
# Max-heap variant of "test_naive_nbest"
|
||||
data = [random.randrange(2000) for i in range(1000)]
|
||||
heap = []
|
||||
for item in data:
|
||||
self.module.heappush_max(heap, item)
|
||||
if len(heap) > 10:
|
||||
self.module.heappop_max(heap)
|
||||
heap.sort()
|
||||
expected = sorted(data)[:10]
|
||||
self.assertEqual(heap, expected)
|
||||
|
||||
def heapiter_max(self, heap):
|
||||
# An iterator returning a max-heap's elements, largest-first.
|
||||
try:
|
||||
while 1:
|
||||
yield self.module.heappop_max(heap)
|
||||
except IndexError:
|
||||
pass
|
||||
|
||||
def test_nworst(self):
|
||||
# Max-heap variant of "test_nbest"
|
||||
data = [random.randrange(2000) for i in range(1000)]
|
||||
heap = data[:10]
|
||||
self.module.heapify_max(heap)
|
||||
for item in data[10:]:
|
||||
if item < heap[0]: # this gets rarer the longer we run
|
||||
self.module.heapreplace_max(heap, item)
|
||||
expected = sorted(data, reverse=True)[-10:]
|
||||
self.assertEqual(list(self.heapiter_max(heap)), expected)
|
||||
|
||||
self.assertRaises(TypeError, self.module.heapreplace_max, None)
|
||||
self.assertRaises(TypeError, self.module.heapreplace_max, None, None)
|
||||
self.assertRaises(IndexError, self.module.heapreplace_max, [], None)
|
||||
|
||||
def test_nworst_minheap(self):
|
||||
# Min-heap variant of "test_nbest_maxheap"
|
||||
data = [random.randrange(2000) for i in range(1000)]
|
||||
heap = data[:]
|
||||
self.module.heapify(heap)
|
||||
result = [self.module.heappop(heap) for _ in range(10)]
|
||||
result.reverse()
|
||||
expected = sorted(data, reverse=True)[-10:]
|
||||
self.assertEqual(result, expected)
|
||||
|
||||
def test_nworst_with_pushpop(self):
|
||||
# Max-heap variant of "test_nbest_with_pushpop"
|
||||
data = [random.randrange(2000) for i in range(1000)]
|
||||
heap = data[:10]
|
||||
self.module.heapify_max(heap)
|
||||
for item in data[10:]:
|
||||
self.module.heappushpop_max(heap, item)
|
||||
expected = sorted(data, reverse=True)[-10:]
|
||||
self.assertEqual(list(self.heapiter_max(heap)), expected)
|
||||
self.assertEqual(self.module.heappushpop_max([], 'x'), 'x')
|
||||
|
||||
def test_heappushpop(self):
|
||||
h = []
|
||||
x = self.module.heappushpop(h, 10)
|
||||
|
@ -153,12 +259,31 @@ class TestHeap:
|
|||
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.assertTupleEqual((h, x), ([], 10))
|
||||
|
||||
h = [10]
|
||||
x = self.module.heappushpop_max(h, 10.0)
|
||||
self.assertTupleEqual((h, x), ([10], 10.0))
|
||||
self.assertIsInstance(h[0], int)
|
||||
self.assertIsInstance(x, float)
|
||||
|
||||
h = [10]
|
||||
x = self.module.heappushpop_max(h, 11)
|
||||
self.assertTupleEqual((h, x), ([10], 11))
|
||||
|
||||
h = [10]
|
||||
x = self.module.heappushpop_max(h, 9)
|
||||
self.assertTupleEqual((h, x), ([9], 10))
|
||||
|
||||
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
|
||||
|
@ -175,6 +300,20 @@ class TestHeap:
|
|||
heap_sorted = [self.module.heappop(heap) for i in range(size)]
|
||||
self.assertEqual(heap_sorted, sorted(data))
|
||||
|
||||
def test_heapsort_max(self):
|
||||
for trial in range(100):
|
||||
size = random.randrange(50)
|
||||
data = [random.randrange(25) for i in range(size)]
|
||||
if trial & 1: # Half of the time, use heapify_max
|
||||
heap = data[:]
|
||||
self.module.heapify_max(heap)
|
||||
else: # The rest of the time, use heappush_max
|
||||
heap = []
|
||||
for item in data:
|
||||
self.module.heappush_max(heap, item)
|
||||
heap_sorted = [self.module.heappop_max(heap) for i in range(size)]
|
||||
self.assertEqual(heap_sorted, sorted(data, reverse=True))
|
||||
|
||||
def test_merge(self):
|
||||
inputs = []
|
||||
for i in range(random.randrange(25)):
|
||||
|
@ -377,16 +516,20 @@ class SideEffectLT:
|
|||
class TestErrorHandling:
|
||||
|
||||
def test_non_sequence(self):
|
||||
for f in (self.module.heapify, self.module.heappop):
|
||||
for f in (self.module.heapify, self.module.heappop,
|
||||
self.module.heapify_max, self.module.heappop_max):
|
||||
self.assertRaises((TypeError, AttributeError), f, 10)
|
||||
for f in (self.module.heappush, self.module.heapreplace,
|
||||
self.module.heappush_max, self.module.heapreplace_max,
|
||||
self.module.nlargest, self.module.nsmallest):
|
||||
self.assertRaises((TypeError, AttributeError), f, 10, 10)
|
||||
|
||||
def test_len_only(self):
|
||||
for f in (self.module.heapify, self.module.heappop):
|
||||
for f in (self.module.heapify, self.module.heappop,
|
||||
self.module.heapify_max, self.module.heappop_max):
|
||||
self.assertRaises((TypeError, AttributeError), f, LenOnly())
|
||||
for f in (self.module.heappush, self.module.heapreplace):
|
||||
for f in (self.module.heappush, self.module.heapreplace,
|
||||
self.module.heappush_max, self.module.heapreplace_max):
|
||||
self.assertRaises((TypeError, AttributeError), f, LenOnly(), 10)
|
||||
for f in (self.module.nlargest, self.module.nsmallest):
|
||||
self.assertRaises(TypeError, f, 2, LenOnly())
|
||||
|
@ -395,7 +538,8 @@ class TestErrorHandling:
|
|||
seq = [CmpErr(), CmpErr(), CmpErr()]
|
||||
for f in (self.module.heapify, self.module.heappop):
|
||||
self.assertRaises(ZeroDivisionError, f, seq)
|
||||
for f in (self.module.heappush, self.module.heapreplace):
|
||||
for f in (self.module.heappush, self.module.heapreplace,
|
||||
self.module.heappush_max, self.module.heapreplace_max):
|
||||
self.assertRaises(ZeroDivisionError, f, seq, 10)
|
||||
for f in (self.module.nlargest, self.module.nsmallest):
|
||||
self.assertRaises(ZeroDivisionError, f, 2, seq)
|
||||
|
@ -403,6 +547,8 @@ class TestErrorHandling:
|
|||
def test_arg_parsing(self):
|
||||
for f in (self.module.heapify, self.module.heappop,
|
||||
self.module.heappush, self.module.heapreplace,
|
||||
self.module.heapify_max, self.module.heappop_max,
|
||||
self.module.heappush_max, self.module.heapreplace_max,
|
||||
self.module.nlargest, self.module.nsmallest):
|
||||
self.assertRaises((TypeError, AttributeError), f, 10)
|
||||
|
||||
|
@ -424,6 +570,10 @@ class TestErrorHandling:
|
|||
# Python version raises IndexError, C version RuntimeError
|
||||
with self.assertRaises((IndexError, RuntimeError)):
|
||||
self.module.heappush(heap, SideEffectLT(5, heap))
|
||||
heap = []
|
||||
heap.extend(SideEffectLT(i, heap) for i in range(200))
|
||||
with self.assertRaises((IndexError, RuntimeError)):
|
||||
self.module.heappush_max(heap, SideEffectLT(5, heap))
|
||||
|
||||
def test_heappop_mutating_heap(self):
|
||||
heap = []
|
||||
|
@ -431,8 +581,12 @@ class TestErrorHandling:
|
|||
# Python version raises IndexError, C version RuntimeError
|
||||
with self.assertRaises((IndexError, RuntimeError)):
|
||||
self.module.heappop(heap)
|
||||
heap = []
|
||||
heap.extend(SideEffectLT(i, heap) for i in range(200))
|
||||
with self.assertRaises((IndexError, RuntimeError)):
|
||||
self.module.heappop_max(heap)
|
||||
|
||||
def test_comparison_operator_modifiying_heap(self):
|
||||
def test_comparison_operator_modifying_heap(self):
|
||||
# See bpo-39421: Strong references need to be taken
|
||||
# when comparing objects as they can alter the heap
|
||||
class EvilClass(int):
|
||||
|
@ -444,7 +598,7 @@ class TestErrorHandling:
|
|||
self.module.heappush(heap, EvilClass(0))
|
||||
self.assertRaises(IndexError, self.module.heappushpop, heap, 1)
|
||||
|
||||
def test_comparison_operator_modifiying_heap_two_heaps(self):
|
||||
def test_comparison_operator_modifying_heap_two_heaps(self):
|
||||
|
||||
class h(int):
|
||||
def __lt__(self, o):
|
||||
|
@ -464,6 +618,17 @@ class TestErrorHandling:
|
|||
self.assertRaises((IndexError, RuntimeError), self.module.heappush, list1, g(1))
|
||||
self.assertRaises((IndexError, RuntimeError), self.module.heappush, list2, h(1))
|
||||
|
||||
list1, list2 = [], []
|
||||
|
||||
self.module.heappush_max(list1, h(0))
|
||||
self.module.heappush_max(list2, g(0))
|
||||
self.module.heappush_max(list1, g(1))
|
||||
self.module.heappush_max(list2, h(1))
|
||||
|
||||
self.assertRaises((IndexError, RuntimeError), self.module.heappush_max, list1, g(1))
|
||||
self.assertRaises((IndexError, RuntimeError), self.module.heappush_max, list2, h(1))
|
||||
|
||||
|
||||
class TestErrorHandlingPython(TestErrorHandling, TestCase):
|
||||
module = py_heapq
|
||||
|
||||
|
|
|
@ -0,0 +1,5 @@
|
|||
Make :mod:`heapq` max-heap functions :func:`heapq.heapify_max`, :func:`heapq.heappush_max`,
|
||||
:func:`heapq.heappop_max`, and :func:`heapq.heapreplace_max` public.
|
||||
Previous underscored naming is kept for backwards compatibility.
|
||||
Additionally, the missing function :func:`heapq.heappushpop_max` has been added
|
||||
to both the C and Python implementations.
|
|
@ -480,9 +480,33 @@ siftup_max(PyListObject *heap, Py_ssize_t pos)
|
|||
return siftdown_max(heap, startpos, pos);
|
||||
}
|
||||
|
||||
/*[clinic input]
|
||||
_heapq.heappush_max
|
||||
|
||||
heap: object(subclass_of='&PyList_Type')
|
||||
item: object
|
||||
/
|
||||
|
||||
Push item onto max heap, maintaining the heap invariant.
|
||||
[clinic start generated code]*/
|
||||
|
||||
static PyObject *
|
||||
_heapq_heappush_max_impl(PyObject *module, PyObject *heap, PyObject *item)
|
||||
/*[clinic end generated code: output=c869d5f9deb08277 input=4743d7db137b6e2b]*/
|
||||
{
|
||||
if (PyList_Append(heap, item)) {
|
||||
return NULL;
|
||||
}
|
||||
|
||||
if (siftdown_max((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-1)) {
|
||||
return NULL;
|
||||
}
|
||||
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
/*[clinic input]
|
||||
_heapq._heappop_max
|
||||
_heapq.heappop_max
|
||||
|
||||
heap: object(subclass_of='&PyList_Type')
|
||||
/
|
||||
|
@ -491,14 +515,14 @@ Maxheap variant of heappop.
|
|||
[clinic start generated code]*/
|
||||
|
||||
static PyObject *
|
||||
_heapq__heappop_max_impl(PyObject *module, PyObject *heap)
|
||||
/*[clinic end generated code: output=9e77aadd4e6a8760 input=362c06e1c7484793]*/
|
||||
_heapq_heappop_max_impl(PyObject *module, PyObject *heap)
|
||||
/*[clinic end generated code: output=2f051195ab404b77 input=e62b14016a5a26de]*/
|
||||
{
|
||||
return heappop_internal(heap, siftup_max);
|
||||
}
|
||||
|
||||
/*[clinic input]
|
||||
_heapq._heapreplace_max
|
||||
_heapq.heapreplace_max
|
||||
|
||||
heap: object(subclass_of='&PyList_Type')
|
||||
item: object
|
||||
|
@ -508,15 +532,14 @@ Maxheap variant of heapreplace.
|
|||
[clinic start generated code]*/
|
||||
|
||||
static PyObject *
|
||||
_heapq__heapreplace_max_impl(PyObject *module, PyObject *heap,
|
||||
PyObject *item)
|
||||
/*[clinic end generated code: output=8ad7545e4a5e8adb input=f2dd27cbadb948d7]*/
|
||||
_heapq_heapreplace_max_impl(PyObject *module, PyObject *heap, PyObject *item)
|
||||
/*[clinic end generated code: output=8770778b5a9cbe9b input=21a3d28d757c881c]*/
|
||||
{
|
||||
return heapreplace_internal(heap, item, siftup_max);
|
||||
}
|
||||
|
||||
/*[clinic input]
|
||||
_heapq._heapify_max
|
||||
_heapq.heapify_max
|
||||
|
||||
heap: object(subclass_of='&PyList_Type')
|
||||
/
|
||||
|
@ -525,21 +548,74 @@ Maxheap variant of heapify.
|
|||
[clinic start generated code]*/
|
||||
|
||||
static PyObject *
|
||||
_heapq__heapify_max_impl(PyObject *module, PyObject *heap)
|
||||
/*[clinic end generated code: output=2cb028beb4a8b65e input=c1f765ee69f124b8]*/
|
||||
_heapq_heapify_max_impl(PyObject *module, PyObject *heap)
|
||||
/*[clinic end generated code: output=8401af3856529807 input=edda4255728c431e]*/
|
||||
{
|
||||
return heapify_internal(heap, siftup_max);
|
||||
}
|
||||
|
||||
/*[clinic input]
|
||||
_heapq.heappushpop_max
|
||||
|
||||
heap: object(subclass_of='&PyList_Type')
|
||||
item: object
|
||||
/
|
||||
|
||||
Maxheap variant of heappushpop.
|
||||
|
||||
The combined action runs more efficiently than heappush_max() followed by
|
||||
a separate call to heappop_max().
|
||||
[clinic start generated code]*/
|
||||
|
||||
static PyObject *
|
||||
_heapq_heappushpop_max_impl(PyObject *module, PyObject *heap, PyObject *item)
|
||||
/*[clinic end generated code: output=ff0019f0941aca0d input=525a843013cbd6c0]*/
|
||||
{
|
||||
PyObject *returnitem;
|
||||
int cmp;
|
||||
|
||||
if (PyList_GET_SIZE(heap) == 0) {
|
||||
return Py_NewRef(item);
|
||||
}
|
||||
|
||||
PyObject *top = PyList_GET_ITEM(heap, 0);
|
||||
Py_INCREF(top);
|
||||
cmp = PyObject_RichCompareBool(item, top, Py_LT);
|
||||
Py_DECREF(top);
|
||||
if (cmp < 0) {
|
||||
return NULL;
|
||||
}
|
||||
if (cmp == 0) {
|
||||
return Py_NewRef(item);
|
||||
}
|
||||
|
||||
if (PyList_GET_SIZE(heap) == 0) {
|
||||
PyErr_SetString(PyExc_IndexError, "index out of range");
|
||||
return NULL;
|
||||
}
|
||||
|
||||
returnitem = PyList_GET_ITEM(heap, 0);
|
||||
PyList_SET_ITEM(heap, 0, Py_NewRef(item));
|
||||
if (siftup_max((PyListObject *)heap, 0) < 0) {
|
||||
Py_DECREF(returnitem);
|
||||
return NULL;
|
||||
}
|
||||
return returnitem;
|
||||
}
|
||||
|
||||
static PyMethodDef heapq_methods[] = {
|
||||
_HEAPQ_HEAPPUSH_METHODDEF
|
||||
_HEAPQ_HEAPPUSHPOP_METHODDEF
|
||||
_HEAPQ_HEAPPOP_METHODDEF
|
||||
_HEAPQ_HEAPREPLACE_METHODDEF
|
||||
_HEAPQ_HEAPIFY_METHODDEF
|
||||
_HEAPQ__HEAPPOP_MAX_METHODDEF
|
||||
_HEAPQ__HEAPIFY_MAX_METHODDEF
|
||||
_HEAPQ__HEAPREPLACE_MAX_METHODDEF
|
||||
|
||||
_HEAPQ_HEAPPUSH_MAX_METHODDEF
|
||||
_HEAPQ_HEAPPUSHPOP_MAX_METHODDEF
|
||||
_HEAPQ_HEAPPOP_MAX_METHODDEF
|
||||
_HEAPQ_HEAPREPLACE_MAX_METHODDEF
|
||||
_HEAPQ_HEAPIFY_MAX_METHODDEF
|
||||
|
||||
{NULL, NULL} /* sentinel */
|
||||
};
|
||||
|
||||
|
|
170
Modules/clinic/_heapqmodule.c.h
generated
170
Modules/clinic/_heapqmodule.c.h
generated
|
@ -175,96 +175,166 @@ exit:
|
|||
return return_value;
|
||||
}
|
||||
|
||||
PyDoc_STRVAR(_heapq__heappop_max__doc__,
|
||||
"_heappop_max($module, heap, /)\n"
|
||||
PyDoc_STRVAR(_heapq_heappush_max__doc__,
|
||||
"heappush_max($module, heap, item, /)\n"
|
||||
"--\n"
|
||||
"\n"
|
||||
"Maxheap variant of heappop.");
|
||||
"Push item onto max heap, maintaining the heap invariant.");
|
||||
|
||||
#define _HEAPQ__HEAPPOP_MAX_METHODDEF \
|
||||
{"_heappop_max", (PyCFunction)_heapq__heappop_max, METH_O, _heapq__heappop_max__doc__},
|
||||
#define _HEAPQ_HEAPPUSH_MAX_METHODDEF \
|
||||
{"heappush_max", _PyCFunction_CAST(_heapq_heappush_max), METH_FASTCALL, _heapq_heappush_max__doc__},
|
||||
|
||||
static PyObject *
|
||||
_heapq__heappop_max_impl(PyObject *module, PyObject *heap);
|
||||
_heapq_heappush_max_impl(PyObject *module, PyObject *heap, PyObject *item);
|
||||
|
||||
static PyObject *
|
||||
_heapq__heappop_max(PyObject *module, PyObject *arg)
|
||||
{
|
||||
PyObject *return_value = NULL;
|
||||
PyObject *heap;
|
||||
|
||||
if (!PyList_Check(arg)) {
|
||||
_PyArg_BadArgument("_heappop_max", "argument", "list", arg);
|
||||
goto exit;
|
||||
}
|
||||
heap = arg;
|
||||
return_value = _heapq__heappop_max_impl(module, heap);
|
||||
|
||||
exit:
|
||||
return return_value;
|
||||
}
|
||||
|
||||
PyDoc_STRVAR(_heapq__heapreplace_max__doc__,
|
||||
"_heapreplace_max($module, heap, item, /)\n"
|
||||
"--\n"
|
||||
"\n"
|
||||
"Maxheap variant of heapreplace.");
|
||||
|
||||
#define _HEAPQ__HEAPREPLACE_MAX_METHODDEF \
|
||||
{"_heapreplace_max", _PyCFunction_CAST(_heapq__heapreplace_max), METH_FASTCALL, _heapq__heapreplace_max__doc__},
|
||||
|
||||
static PyObject *
|
||||
_heapq__heapreplace_max_impl(PyObject *module, PyObject *heap,
|
||||
PyObject *item);
|
||||
|
||||
static PyObject *
|
||||
_heapq__heapreplace_max(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
|
||||
_heapq_heappush_max(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
|
||||
{
|
||||
PyObject *return_value = NULL;
|
||||
PyObject *heap;
|
||||
PyObject *item;
|
||||
|
||||
if (!_PyArg_CheckPositional("_heapreplace_max", nargs, 2, 2)) {
|
||||
if (!_PyArg_CheckPositional("heappush_max", nargs, 2, 2)) {
|
||||
goto exit;
|
||||
}
|
||||
if (!PyList_Check(args[0])) {
|
||||
_PyArg_BadArgument("_heapreplace_max", "argument 1", "list", args[0]);
|
||||
_PyArg_BadArgument("heappush_max", "argument 1", "list", args[0]);
|
||||
goto exit;
|
||||
}
|
||||
heap = args[0];
|
||||
item = args[1];
|
||||
return_value = _heapq__heapreplace_max_impl(module, heap, item);
|
||||
return_value = _heapq_heappush_max_impl(module, heap, item);
|
||||
|
||||
exit:
|
||||
return return_value;
|
||||
}
|
||||
|
||||
PyDoc_STRVAR(_heapq__heapify_max__doc__,
|
||||
"_heapify_max($module, heap, /)\n"
|
||||
PyDoc_STRVAR(_heapq_heappop_max__doc__,
|
||||
"heappop_max($module, heap, /)\n"
|
||||
"--\n"
|
||||
"\n"
|
||||
"Maxheap variant of heapify.");
|
||||
"Maxheap variant of heappop.");
|
||||
|
||||
#define _HEAPQ__HEAPIFY_MAX_METHODDEF \
|
||||
{"_heapify_max", (PyCFunction)_heapq__heapify_max, METH_O, _heapq__heapify_max__doc__},
|
||||
#define _HEAPQ_HEAPPOP_MAX_METHODDEF \
|
||||
{"heappop_max", (PyCFunction)_heapq_heappop_max, METH_O, _heapq_heappop_max__doc__},
|
||||
|
||||
static PyObject *
|
||||
_heapq__heapify_max_impl(PyObject *module, PyObject *heap);
|
||||
_heapq_heappop_max_impl(PyObject *module, PyObject *heap);
|
||||
|
||||
static PyObject *
|
||||
_heapq__heapify_max(PyObject *module, PyObject *arg)
|
||||
_heapq_heappop_max(PyObject *module, PyObject *arg)
|
||||
{
|
||||
PyObject *return_value = NULL;
|
||||
PyObject *heap;
|
||||
|
||||
if (!PyList_Check(arg)) {
|
||||
_PyArg_BadArgument("_heapify_max", "argument", "list", arg);
|
||||
_PyArg_BadArgument("heappop_max", "argument", "list", arg);
|
||||
goto exit;
|
||||
}
|
||||
heap = arg;
|
||||
return_value = _heapq__heapify_max_impl(module, heap);
|
||||
return_value = _heapq_heappop_max_impl(module, heap);
|
||||
|
||||
exit:
|
||||
return return_value;
|
||||
}
|
||||
/*[clinic end generated code: output=05f2afdf3bc54c9d input=a9049054013a1b77]*/
|
||||
|
||||
PyDoc_STRVAR(_heapq_heapreplace_max__doc__,
|
||||
"heapreplace_max($module, heap, item, /)\n"
|
||||
"--\n"
|
||||
"\n"
|
||||
"Maxheap variant of heapreplace.");
|
||||
|
||||
#define _HEAPQ_HEAPREPLACE_MAX_METHODDEF \
|
||||
{"heapreplace_max", _PyCFunction_CAST(_heapq_heapreplace_max), METH_FASTCALL, _heapq_heapreplace_max__doc__},
|
||||
|
||||
static PyObject *
|
||||
_heapq_heapreplace_max_impl(PyObject *module, PyObject *heap, PyObject *item);
|
||||
|
||||
static PyObject *
|
||||
_heapq_heapreplace_max(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
|
||||
{
|
||||
PyObject *return_value = NULL;
|
||||
PyObject *heap;
|
||||
PyObject *item;
|
||||
|
||||
if (!_PyArg_CheckPositional("heapreplace_max", nargs, 2, 2)) {
|
||||
goto exit;
|
||||
}
|
||||
if (!PyList_Check(args[0])) {
|
||||
_PyArg_BadArgument("heapreplace_max", "argument 1", "list", args[0]);
|
||||
goto exit;
|
||||
}
|
||||
heap = args[0];
|
||||
item = args[1];
|
||||
return_value = _heapq_heapreplace_max_impl(module, heap, item);
|
||||
|
||||
exit:
|
||||
return return_value;
|
||||
}
|
||||
|
||||
PyDoc_STRVAR(_heapq_heapify_max__doc__,
|
||||
"heapify_max($module, heap, /)\n"
|
||||
"--\n"
|
||||
"\n"
|
||||
"Maxheap variant of heapify.");
|
||||
|
||||
#define _HEAPQ_HEAPIFY_MAX_METHODDEF \
|
||||
{"heapify_max", (PyCFunction)_heapq_heapify_max, METH_O, _heapq_heapify_max__doc__},
|
||||
|
||||
static PyObject *
|
||||
_heapq_heapify_max_impl(PyObject *module, PyObject *heap);
|
||||
|
||||
static PyObject *
|
||||
_heapq_heapify_max(PyObject *module, PyObject *arg)
|
||||
{
|
||||
PyObject *return_value = NULL;
|
||||
PyObject *heap;
|
||||
|
||||
if (!PyList_Check(arg)) {
|
||||
_PyArg_BadArgument("heapify_max", "argument", "list", arg);
|
||||
goto exit;
|
||||
}
|
||||
heap = arg;
|
||||
return_value = _heapq_heapify_max_impl(module, heap);
|
||||
|
||||
exit:
|
||||
return return_value;
|
||||
}
|
||||
|
||||
PyDoc_STRVAR(_heapq_heappushpop_max__doc__,
|
||||
"heappushpop_max($module, heap, item, /)\n"
|
||||
"--\n"
|
||||
"\n"
|
||||
"Maxheap variant of heappushpop.\n"
|
||||
"\n"
|
||||
"The combined action runs more efficiently than heappush_max() followed by\n"
|
||||
"a separate call to heappop_max().");
|
||||
|
||||
#define _HEAPQ_HEAPPUSHPOP_MAX_METHODDEF \
|
||||
{"heappushpop_max", _PyCFunction_CAST(_heapq_heappushpop_max), METH_FASTCALL, _heapq_heappushpop_max__doc__},
|
||||
|
||||
static PyObject *
|
||||
_heapq_heappushpop_max_impl(PyObject *module, PyObject *heap, PyObject *item);
|
||||
|
||||
static PyObject *
|
||||
_heapq_heappushpop_max(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
|
||||
{
|
||||
PyObject *return_value = NULL;
|
||||
PyObject *heap;
|
||||
PyObject *item;
|
||||
|
||||
if (!_PyArg_CheckPositional("heappushpop_max", nargs, 2, 2)) {
|
||||
goto exit;
|
||||
}
|
||||
if (!PyList_Check(args[0])) {
|
||||
_PyArg_BadArgument("heappushpop_max", "argument 1", "list", args[0]);
|
||||
goto exit;
|
||||
}
|
||||
heap = args[0];
|
||||
item = args[1];
|
||||
return_value = _heapq_heappushpop_max_impl(module, heap, item);
|
||||
|
||||
exit:
|
||||
return return_value;
|
||||
}
|
||||
/*[clinic end generated code: output=f55d8595ce150c76 input=a9049054013a1b77]*/
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue