Search code examples
pythoniteratorgeneratorpython-itertools

Get multiple individual values from generator in Python


How do i take multiple arbitrary values from different index positions in an iterator?

How to get the n next values of a generator in a list (python) and Get the nth item of a generator in Python describe the use of itertools.islice for taking arbitrary element or continuous subset from an iterator. But what if i want multiple arbitrary elements from different positions in the iterator, where you can't just use islice's step argument?

I'm trying to solve Project Euler's problem 40. I generated a string of concatenated integers with

iteration = (i for i in ''.join(map(str, (i for i in xrange(1,10**6)))))

An now i want to get elements with indexes 1, 10, 100, 1000, 10000, 100000, 1000000 counting from 1. I couldn't use islice here, because every call to next shifts current value to yield to the right. For example

next(islice(iteration, 1, 2)) + next(islice(iteration, 3, 4))

produces '26' instead of '24'.

Update (25.11.12, 4:43 UTC+0):

Thanks for all the suggestions. My current code looks like:

it = (i for i in ''.join(map(str, (i for i in xrange(1,10**6)))))
ds = [int(nth(it, 10**i-10**(i-1)-1)) for i in range(7)]
return product(ds)

The ugly argument for nth is to generate a sequence of 0, 8, 89, 899, 8999 etc.


Solution

  • This is from the "recipes" section of the itertools documentation. It returns the nth element of iterable, consuming it as it goes:

    def nth(iterable, n, default=None):
        "Returns the nth item or a default value"
        return next(islice(iterable, n, None), default)
    

    You can get the 1st, 10th, 100th etc elements by calling it sequentially (noting that the iterator gets consumed, and is zero-indexed):

    first = nth(iteration, 0)
    tenth = nth(iteration, 8)  # since we've already taken one
    hundredth = nth(iteration, 89)  # since we've already taken ten
    # etc
    

    Alternatively, you could use tee and use nth with a different iterator each time. This way you don't have to worry about the fact that your single iterator is getting consumed. On the other hand you might start swallowing memory if your iterators are very long.