Get the nth item of a generator in Python

Question

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.

Louvain

You don't want is in this situation (or many situations at all). is is for comparing identity, not equality. You want ==. This will probably work in this instance, but only by coincidence and implementation detail. – Width

Since I'm using integers, how could it not work? Is it even good practice to expect the index object to implement __eq__ in cases such as this? (This is getting off topic...) – Louvain

Try 1000 is 500 + 500, it will (probably) be False. See, for example, #306813 – Gogh

+1 for this question. It does seem strange that there's not a less verbose way to say "the nth result of gen". – Attachment

Another possibility is zippers --- they handle arbitrary trees, but a list is a tree too. See this implementation github.com/trivio/zipper/blob/master/tests/test_zipper.py – Rann

Use is for is True or is None. For int using == is fine, but not for float where you should compare a difference to a tolerance like < 1e-7. – Nymphalid

Answer 1

one method would be to use itertools.islice

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

Biggs

Answer 2

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.

Billon

Answer 3

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

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

Caviness

What version of Python is this? The above code gives me the error AttributeError: 'itertools.islice' object has no attribute 'next' in 3.3. – Attachment

In Python 3x, change next to __next__(), i.e., islice(count, n, n=1).__next__() – Mooned

So it's better to use next(islice(count(), n, n+1)). – Pomposity

I think you can get rid of the upper bound, i.e. next(islice(count(), n, None)). – Marasmus

Answer 4

I'd argue against the temptation to treat generators like lists. The simple but naive approach is the simple one-liner:

gen = (i for i in range(10))
list(gen)[3]

But remember, generators aren't like lists. They don't store their intermediate results anywhere, so you can't go backwards. I'll demonstrate the problem with a simple example in the python repl:

>>> gen = (i for i in range(10))
>>> list(gen)[3]
3
>>> list(gen)[3]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: list index out of range

Once you start going through a generator to get the nth value in the sequence, the generator is now in a different state, and attempting to get the nth value again will return you a different result, which is likely to result in a bug in your code.

Let's take a look at another example, based on the code from the question.

One would initially expect the following to print 4 twice.

gen = (i for i in range(10))
index = 4
for i, v in enumerate(gen):
    if i == index:
        answer = v
        break
print(answer)
for i, v in enumerate(gen):
    if i == index:
        answer = v
        break
print(answer)

but type this into the repl and you get:

>>> gen = (i for i in range(10))
>>> index = 4
>>> for i, v in enumerate(gen):
...     if i == index:
...             answer = v
...             break
... 
>>> print(answer)
4
>>> for i, v in enumerate(gen):
...     if i == index:
...             answer = v
...             break
... 
>>> print(answer)
9

Good luck tracing that bug down.

---

As pointed out, if the generator is infinitely long, you can't even convert it to a list. The expression list(gen) will never finish.

There is a way you could put a lazily evaluated caching wrapper around an infinite generator to make it look like an infinitely long list you could index into at will, but that deserves its own question and answer, and would have major performance implications.

Pru

What if the generator is infinite? – Rattat

This should be higher up, as doing this comes at a significant cost in terms of time. Thank you for pointing this out. – Sedulity

Answer 5

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]

Abner

Answer 6

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

Patti

I like it a lot, this doesn't import itertools or pretends to create a "slice", instead simply "jumps" through the iterator and returns the next() value. – Banff

Check out my take on this, with some criticisms =) – Banff

@KlasŠ. you misread my solution. As "next(gen) and False" is always false, it never adds anything to the list. It only creates an empty list, make some kind of empty loop then discard the still empty list. It means to be memory efficient. – Patti

Right you are, totally missed that, my bad =) – Banff

Answer 7

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

Worms

Answer 8

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 =)

Innovation

Answer 9

Perhaps you should elaborate more on a actual use case.

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

Mekka

I editted xrange(10) to (i for i in xrange(10)). Turns out this syntax works for xrange since it's not really a generator... – Louvain

xrange predates generators, and returns an xrange object, which actually implements the full sequence protocol. – Width

Answer 10

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.

Quintinquintina