[python] How to take the first N items from a generator or list?

With I would

var top5 = array.Take(5);

How to do this with Python?

This question is related to python list generator

The answer is


Do you mean the first N items, or the N largest items?

If you want the first:

top5 = sequence[:5]

This also works for the largest N items, assuming that your sequence is sorted in descending order. (Your LINQ example seems to assume this as well.)

If you want the largest, and it isn't sorted, the most obvious solution is to sort it first:

l = list(sequence)
l.sort(reverse=True)
top5 = l[:5]

For a more performant solution, use a min-heap (thanks Thijs):

import heapq
top5 = heapq.nlargest(5, sequence)

import itertools

top5 = itertools.islice(array, 5)

The answer for how to do this can be found here

>>> generator = (i for i in xrange(10))
>>> list(next(generator) for _ in range(4))
[0, 1, 2, 3]
>>> list(next(generator) for _ in range(4))
[4, 5, 6, 7]
>>> list(next(generator) for _ in range(4))
[8, 9]

Notice that the last call asks for the next 4 when only 2 are remaining. The use of the list() instead of [] is what gets the comprehension to terminate on the StopIteration exception that is thrown by next().


With itertools you will obtain another generator object so in most of the cases you will need another step the take the first N elements (N). There are at least two simpler solutions (a little bit less efficient in terms of performance but very handy) to get the elements ready to use from a generator:

Using list comprehension:

first_N_element=[generator.next() for i in range(N)]

Otherwise:

first_N_element=list(generator)[:N]

Where N is the number of elements you want to take (e.g. N=5 for the first five elements).


In my taste, it's also very concise to combine zip() with xrange(n) (or range(n) in Python3), which works nice on generators as well and seems to be more flexible for changes in general.

# Option #1: taking the first n elements as a list
[x for _, x in zip(xrange(n), generator)]

# Option #2, using 'next()' and taking care for 'StopIteration'
[next(generator) for _ in xrange(n)]

# Option #3: taking the first n elements as a new generator
(x for _, x in zip(xrange(n), generator))

# Option #4: yielding them by simply preparing a function
# (but take care for 'StopIteration')
def top_n(n, generator):
    for _ in xrange(n): yield next(generator)

@Shaikovsky's answer is excellent (…and heavily edited since I posted this answer), but I wanted to clarify a couple of points.

[next(generator) for _ in range(n)]

This is the most simple approach, but throws StopIteration if the generator is prematurely exhausted.


On the other hand, the following approaches return up to n items which is preferable in many circumstances:

List: [x for _, x in zip(range(n), records)]

Generator: (x for _, x in zip(range(n), records))


This should work

top5 = array[:5] 

Examples related to python

programming a servo thru a barometer Is there a way to view two blocks of code from the same file simultaneously in Sublime Text? python variable NameError Why my regexp for hyphenated words doesn't work? Comparing a variable with a string python not working when redirecting from bash script is it possible to add colors to python output? Get Public URL for File - Google Cloud Storage - App Engine (Python) Real time face detection OpenCV, Python xlrd.biffh.XLRDError: Excel xlsx file; not supported Could not load dynamic library 'cudart64_101.dll' on tensorflow CPU-only installation

Examples related to list

Convert List to Pandas Dataframe Column Python find elements in one list that are not in the other Sorting a list with stream.sorted() in Java Python Loop: List Index Out of Range How to combine two lists in R How do I multiply each element in a list by a number? Save a list to a .txt file The most efficient way to remove first N elements in a list? TypeError: list indices must be integers or slices, not str Parse JSON String into List<string>

Examples related to generator

Gradient text color Is there a mechanism to loop x times in ES6 (ECMAScript 6) without mutable variables? Convert generator object to list for debugging How can I generate a random number in a certain range? Display SQL query results in php What does yield mean in PHP? How does C#'s random number generator work? How to len(generator()) How to take the first N items from a generator or list? How to pick just one item from a generator?