[Python-checkins] python/dist/src/Lib heapq.py,1.27,1.28

Thu Dec 2 09:59:16 CET 2004

Update of /cvsroot/python/python/dist/src/Lib
In directory sc8-pr-cvs1.sourceforge.net:/tmp/cvs-serv20862/Lib
Modified Files:
	heapq.py 
Log Message:
Add key= argument to heapq.nsmallest() and heapq.nlargest().
Index: heapq.py
===================================================================
RCS file: /cvsroot/python/python/dist/src/Lib/heapq.py,v
retrieving revision 1.27
retrieving revision 1.28
diff -u -d -r1.27 -r1.28

--- heapq.py	29 Nov 2004 05:54:47 -0000	1.27
+++ heapq.py	2 Dec 2004 08:59:13 -0000	1.28
@@ -129,7 +129,8 @@
 __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 @@
 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 = []
