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

Question

With linq I would

var top5 = array.Take(5);

How to do this with Python?

Wetnurse

@ThorSummoner I guess OP assumed there's a unified way to do it in Python. – Prefigure

Just to note: I flipped the duplicate closure of this question w/r/t Fetch first 10 results from a list in Python because this question is sort of asking two things, but that question is only asking one thing, so I hope it's easier to follow, especially for beginners. – Prefigure

Yeah, Checkout the Python Data Model and more technically, Python's collections.abc -- which unfortunately require a little interpretation, but we can see that Generators have close, __iter__, __next__ methods, and Lists (which are likely most similar to MutableSequences) have one overlapping method: __iter__. therefore iter([0, 1]) and def fib():... iter(fib()) will provide a common interface to either a generator or list, in the form of a generator, see: islice – Caravel

Answer 1

Slicing a list

top5 = array[:5]

• To slice a list, there's a simple syntax: array[start:stop:step]
• You can omit any parameter. These are all valid: array[start:], array[:stop], array[::step]

Slicing a generator

import itertools
top5 = itertools.islice(my_list, 5) # grab the first five elements

• You can't slice a generator directly in Python. itertools.islice() will wrap an object in a new slicing generator using the syntax itertools.islice(generator, start, stop, step)

• Remember, slicing a generator will exhaust it partially. If you want to keep the entire generator intact, perhaps turn it into a tuple or list first, like: result = tuple(generator)

Airlie

Also note that itertools.islice will return a generator. – Ruffianism

"If you want to keep the entire generator intact, perhaps turn it into a tuple or list first" -> won't that exhaust the generator fully, in the process of building up the tuple / list? – Bedfast

@Bedfast yes, but then you have a new data structure (tuple/list) that you can iterate over as much as you like – Arouse

To create copies of the generator before exhausting it, you can also use itertools.tee, e.g.: generator, another_copy = itertools.tee(generator) – Sean

Note: which slice gets which elements is determined by the order in which the slices are exhausted not in which they are created. import itertools as it;r=(i for i in range(10));s1=itt.islice(r, 5);s2=itt.islice(r, 5);l2=list(s2);l1=list(s1) ends with l1==[5,6,7,8,9] and l2==[0,1,2,3,4] – Slater

Answer 2

import itertools

top5 = itertools.islice(array, 5)

Wetnurse

This also has the nice property of returning the entire array when you have None in place of 5. – Stiletto

and if you want to take the five that follows each time you can use: iter(array) instead of array. – Hobbledehoy

note that if your generator exhausts this will not make an error, you will get a many elements as the generator had left, less than your request size. – Caravel

This is the approach used in the following: Itertools recipes def take(n, iterable): return list(islice(iterable, n)) – Cardenas

Answer 3

@Shaikovsky's answer is excellent, 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))

Opiate

Could those few people downvoting this answer please explain why? – Opiate

def take(num,iterable): return([elem for _ , elem in zip(range(num), iterable)]) – Colossus

Above code: Loop over an iterable which could be a generator or list and return up to n elements from iterable. In case n is greater or equal to number of items existing in iterable then return all elements in iterable. – Colossus

For a list x=[1,2,3,4,5,6], x[:20] also return only the 6 elements in x. Guess x[:N] return first N elements of x, if N > len(x), it will return x. python 3.6. – Vulcan

This is the most efficient. Because this doesn't process the full list. – Ubangishari

[next(generator, None) for _ in range(n)] if you don't mind the None – Hogg

Answer 4

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)

Dun

Answer 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().

Microvolt

Caveats on some python versions – Rizzio

to clarify Tom's comment: in python 3.7 you'll get a RuntimeError (the link is definitely worth the read though!) – Creditor

Answer 6

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)

Gusty

wouldn't the smaller come first? – Wetnurse

use sequence instead of iterable. Some iterables do not support indexing. Every sequence is an iterable, but some iterables are not sequences. – Stearic

Note nlargest takes any iterable, not only sequences. – Intermittent

Answer 7

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. 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_elements = [generator.next() for i in range(n)]

Otherwise:

first_n_elements = list(generator)[:n]

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

Ridglee

Answer 8

This should work

top5 = array[:5] 

Keene

@JoshWolff I didn't downvote this answer, but it's likely because this approach will not work with generators, unless they define __getitem__(). Try running itertools.count()[:5] or (x for x in range(10))[:5], for instance, and see the error messages. The answer is, however, idiomatic for lists. – Cook