Get the nth item of a generator in Python

user_73398 0
Viewed 50249

Is there a more syntactically concise way of writing the following?

gen = (i for i in xrange(10))
index = 5
for i, v in enumerate(gen):
    if i is index:
        return v

It seems almost natural that a generator should have a gen[index] expression, that acts as a list, but is functionally identical to the above code.

generator python
10 Answers

one method would be to use itertools.islice

>>> gen = (x for x in range(10))
>>> index = 5
>>> next(itertools.islice(gen, index, None))
5
user_736810

You could do this, using count as an example generator:

from itertools import islice, count
next(islice(count(), n, n+1))
user_619740

I think the best way is : 

next(x for i,x in enumerate(it) if i==n)

(where it is your iterator and n is the index)

It doesn't require you to add an import (like the solutions using itertools) nor to load all the elements of the iterator in memory at once (like the solutions using list).

Note 1: this version throws a StopIteration error if your iterator has less than n items. If you want to get None instead, you can use :

next((x for i,x in enumerate(it) if i==n), None)

Note 2: There are no brackets inside the call to next. This is not a list comprehension, but a generator comprehension, that does not consume the original iterator further than its nth element.

user_35793090

If n is known at authoring-time, you can use destructuring. e.g. to get the 3rd item:

>>> [_, _, third, *rest] = range(10)
>>> third
2
>>> rest
[3, 4, 5, 6, 7, 8, 9]
user_173280

My solution would be:

[_ for _ in range(n-1) if next(gen) and False ]
return next(gen)

As next(gen) and False is always false, the list comprehension does nothing but executing next(gen) n-1 times.

In my testing, it is as fast as using itertools.islice

user_36551040

The first thing that came to my mind was:

gen = (i for i in xrange(10))
index = 5

for i, v in zip(range(index), gen): pass

return v
user_8986490

Building up on @Madlozoz answer, but with mighty walrus operator:

>>> gen = (x ** 2 for x in itertools.count())
>>> [v := next(gen) for _ in range(10)]
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
>>> v
81

The thing I don't like about this (unlike Madlozoz') solution is that potentially huge list is constructed just to be discarded immediately.

A simple for loop is another way:

# `gen` continues from the previous snippet
>>> for _ in range(3):
...     v = next(gen)
...
>>> print(v)
144

At a cost of extra line we can save some process ticks on assignments as well and spend them on wrapper class:

class IterIndexer:

    def __init__(self, iter_):
        self.iter = iter_

    def __getitem__(self, i):
        for _ in range(i - 1):
            next(self.iter)
        return next(self.iter)


gen = (x ** 2 for x in itertools.count())
gen = IterIndexer(gen)
print(gen[14])

169

It'd be even cooler to wrap it properly, so that you can use wrapper instance for everything instead of original generator or iterator, but that's another question =)

user_92885800

Perhaps you should elaborate more on a actual use case.

>>> gen = xrange(10)
>>> ind=5 
>>> gen[ind]
5
user_1315270

Best to use is : example : 

a = gen values ('a','c','d','e')

so the answer will be : 

a = list(a) -> this will convert the generator to a list (it will store in memory)

then when you want to go specific index you will : 

a[INDEX] -> and you will able to get the value its holds

if you want to know only the count or to do operations that not required store in memory best practice will be : a = sum(1 in i in a) -> this will count the number of objects you have

hope i made it more simple.

user_44798610
Related