SF patch #969791: Add nlargest() and nsmallest() to heapq.
diff --git a/Lib/heapq.py b/Lib/heapq.py
index 3eb69d8..d1aad98 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -30,7 +30,7 @@
maintains the heap invariant!
"""
-# Original code by Kevin O'Connor, augmented by Tim Peters
+# Original code by Kevin O'Connor, augmented by Tim Peters and Raymond Hettinger
__about__ = """Heap queues
@@ -126,7 +126,10 @@
From all times, sorting has always been a Great Art! :-)
"""
-__all__ = ['heappush', 'heappop', 'heapify', 'heapreplace']
+__all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'nlargest',
+ 'nsmallest']
+
+from itertools import islice, repeat
def heappush(heap, item):
"""Push item onto heap, maintaining the heap invariant."""
@@ -168,6 +171,35 @@
for i in reversed(xrange(n//2)):
_siftup(x, i)
+def nlargest(iterable, n):
+ """Find the n largest elements in a dataset.
+
+ Equivalent to: sorted(iterable, reverse=True)[:n]
+ """
+ it = iter(iterable)
+ result = list(islice(it, n))
+ if not result:
+ return result
+ heapify(result)
+ _heapreplace = heapreplace
+ sol = result[0] # sol --> smallest of the nlargest
+ for elem in it:
+ if elem <= sol:
+ continue
+ _heapreplace(result, elem)
+ sol = result[0]
+ result.sort(reverse=True)
+ return result
+
+def nsmallest(iterable, n):
+ """Find the n smallest elements in a dataset.
+
+ Equivalent to: sorted(iterable)[:n]
+ """
+ h = list(iterable)
+ heapify(h)
+ return map(heappop, repeat(h, min(n, len(h))))
+
# '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.