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

rhettinger · rhettinger · commit 4901a1f267e9 · 2004-12-02T08:59:14.000Z
diff --git a/Doc/lib/libheapq.tex b/Doc/lib/libheapq.tex
@@ -90,16 +90,24 @@ \section{\module{heapq} ---
 
 The module also offers two general purpose functions based on heaps.
 
-\begin{funcdesc}{nlargest}{n, iterable}
+\begin{funcdesc}{nlargest}{n, iterable\optional{, key}}
 Return a list with the \var{n} largest elements from the dataset defined
-by \var{iterable}. Equivalent to:  \code{sorted(iterable, reverse=True)[:n]}
-\versionadded{2.4}              
+by \var{iterable}.  \var{key}, if provided, specifies a function of one
+argument that is used to extract a comparison key from each element
+in the iterable:  \samp{\var{key}=\function{str.lower}}
+Equivalent to:  \samp{sorted(iterable, key=key, reverse=True)[:n]}
+\versionadded{2.4}
+\versionchanged[Added the optional \var{key} argument]{2.5}
 \end{funcdesc}
 
-\begin{funcdesc}{nsmallest}{n, iterable}
+\begin{funcdesc}{nsmallest}{n, iterable\optional{, key}}
 Return a list with the \var{n} smallest elements from the dataset defined
-by \var{iterable}. Equivalent to:  \code{sorted(iterable)[:n]}
-\versionadded{2.4}              
+by \var{iterable}.  \var{key}, if provided, specifies a function of one
+argument that is used to extract a comparison key from each element
+in the iterable:  \samp{\var{key}=\function{str.lower}}
+Equivalent to:  \samp{sorted(iterable, key=key)[:n]}
+\versionadded{2.4}
+\versionchanged[Added the optional \var{key} argument]{2.5}
 \end{funcdesc}
 
 Both functions perform best for smaller values of \var{n}.  For larger
diff --git a/Lib/heapq.py b/Lib/heapq.py
@@ -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 @@ def _siftup(heap, pos):
 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 = []
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
@@ -105,13 +105,19 @@ def test_heapsort(self):
 
     def test_nsmallest(self):
         data = [random.randrange(2000) for i in range(1000)]
+        f = lambda x:  x * 547 % 2000
         for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
             self.assertEqual(nsmallest(n, data), sorted(data)[:n])
+            self.assertEqual(nsmallest(n, data, key=f),
+                             sorted(data, key=f)[:n])
 
-    def test_largest(self):
+    def test_nlargest(self):
         data = [random.randrange(2000) for i in range(1000)]
+        f = lambda x:  x * 547 % 2000
         for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
             self.assertEqual(nlargest(n, data), sorted(data, reverse=True)[:n])
+            self.assertEqual(nlargest(n, data, key=f),
+                             sorted(data, key=f, reverse=True)[:n])
 
 
 #==============================================================================
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -4,6 +4,46 @@ Python News
 
 (editors: check NEWS.help for information about editing NEWS using ReST.)
 
+What's New in Python 2.5 alpha 1?
+=================================
+
+Core and builtins
+-----------------
+
+
+Extension Modules
+-----------------
+
+
+Library
+-------
+
+- heapq.nsmallest() and heapq.nlargest() now support key= arguments with
+  the same meaning as for list.sort().
+
+
+Build
+-----
+
+
+C API
+-----
+
+
+Tests
+-----
+
+
+Mac
+---
+
+
+
+Tools/Demos
+-----------
+
+
+
 What's New in Python 2.4 final?
 ===============================
 
