Improve the memory performance and speed of heapq.nsmallest() by using
an alternate algorithm when the number of selected items is small
relative to the full iterable.
diff --git a/Lib/heapq.py b/Lib/heapq.py
index d1aad98..65f4155 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -130,6 +130,7 @@
'nsmallest']
from itertools import islice, repeat
+import bisect
def heappush(heap, item):
"""Push item onto heap, maintaining the heap invariant."""
@@ -196,6 +197,28 @@
Equivalent to: sorted(iterable)[:n]
"""
+ if hasattr(iterable, '__len__') and n * 10 <= len(iterable):
+ # For smaller values of n, the bisect method is faster than a minheap.
+ # It is also memory efficient, consuming only n elements of space.
+ it = iter(iterable)
+ result = sorted(islice(it, 0, n))
+ if not result:
+ return result
+ insort = bisect.insort
+ pop = result.pop
+ los = result[-1] # los --> Largest of the nsmallest
+ for elem in it:
+ if los <= elem:
+ continue
+ insort(result, elem)
+ pop()
+ los = result[-1]
+ return result
+ # An alternative approach manifests the whole iterable in memory but
+ # saves comparisons by heapifying all at once. Also, saves time
+ # over bisect.insort() which has O(n) data movement time for every
+ # insertion. Finding the n smallest of an m length iterable requires
+ # O(m) + O(n log m) comparisons.
h = list(iterable)
heapify(h)
return map(heappop, repeat(h, min(n, len(h))))
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 944b17d..1cdaabe 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -92,6 +92,7 @@
def test_nsmallest(self):
data = [random.randrange(2000) for i in range(1000)]
self.assertEqual(nsmallest(data, 400), sorted(data)[:400])
+ self.assertEqual(nsmallest(data, 50), sorted(data)[:50])
def test_largest(self):
data = [random.randrange(2000) for i in range(1000)]