mirrors/cpython
1
0
Fork 0
mirror of https://github.com/python/cpython.git synced 2025-08-31 05:58:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 7

Add key= argument to heapq.nsmallest() and heapq.nlargest().

Browse source
This commit is contained in:
Raymond Hettinger 2004-12-02 08:59:14 +00:00
parent de7b99045d
commit 4901a1f267
4 changed files with 90 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

30
Lib/heapq.py
View file

@ -129,7 +129,8 @@ From all times, sorting has always been a Great Art! :-)
__all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'nlargest',
'nsmallest']
from itertools import islice, repeat
from itertools import islice, repeat, count, imap, izip, tee
from operator import itemgetter
import bisect
def heappush(heap, item):
@ -307,6 +308,33 @@ try:
except ImportError:
pass
# Extend the implementations of nsmallest and nlargest to use a key= argument
_nsmallest = nsmallest
def nsmallest(n, iterable, key=None):
"""Find the n smallest elements in a dataset.
Equivalent to: sorted(iterable, key=key)[:n]
"""
if key is None:
return _nsmallest(n, iterable)
in1, in2 = tee(iterable)
it = izip(imap(key, in1), count(), in2) # decorate
result = _nsmallest(n, it)
return map(itemgetter(2), result) # undecorate
_nlargest = nlargest
def nlargest(n, iterable, key=None):
"""Find the n largest elements in a dataset.
Equivalent to: sorted(iterable, key=key, reverse=True)[:n]
"""
if key is None:
return _nlargest(n, iterable)
in1, in2 = tee(iterable)
it = izip(imap(key, in1), count(), in2) # decorate
result = _nlargest(n, it)
return map(itemgetter(2), result) # undecorate
if __name__ == "__main__":
# Simple sanity test
heap = []