# Retrieving the nth term of an infinite stream

I recently learned how generators can be infinite in Python. For example, the infinite sequence can be the triangular numbers: $$1, 3, 6, 10,...$$

def triangle_stream():
""" Infinite stream of triangle numbers. """
x = step = 1
yield x
while True:
step += 1
x += step
yield x


I wrote a function that returns the nth term in an infinite sequence.

def nth_term(generator, n):
""" Returns the nth term of an infinite generator. """
for i, x in enumerate(generator()):
if i == n - 1:
return x


As an example:

>>> nth_term(triangle_stream, 10)
55


But this nth_term function also works:

def nth_term(generator, n):
""" Returns the nth term of an infinite generator. """
x = generator()
for i in range(n - 1):
next(x)
return next(x)


Which of these two nth_term functions is more efficient?

• Have you done performance testing? I would guess that the second will be faster because you remove an if statement, but the best way to find out would be to test. Nov 3, 2016 at 4:11
• I'm using repl.it, I'm not sure if performance testing is possible on that Nov 3, 2016 at 4:22
• what about timeit? Nov 3, 2016 at 4:54

Two modifications that I would make right away would be:

1. to use iterables instead of a generator function as parameter to allow for more genericity: this way I could use nth_term using either triangle_stream() or itertools.cycle('ABCD') as parameters;
2. to index elements in base 0 instead of base 1, so nth_term(triangle_stream(), 0) would return 1 instead of looping indefinitely and nth_term(triangle_stream(), 1) would return 3 instead of 1.

Now, making the function accept any iterable, I would use something along the lines of the second version this will be closer to my final proposition that will raise StopIteration if there is not enough elements instead of silently returning None ("Errors should never pass silently."). But applying both suggestions would lead to:

def nth_term(iterable, n):
"""Returns the nth term of any iterable."""
for i, x in enumerate(iterable):
if i == n:
return x


or

def nth_term(iterable, n):
"""Returns the nth term of any iterable."""
x = iter(iterable)
for i in range(n):
next(x)
return next(x)


Now, instead of manually calling next on the iterator in the second version, I would ask zip to do it for me. Chances are that the execution will spend less time in the Python interpreter and more in the C part of zip and so be faster:

def nth_term(iterable, n):
"""Returns the nth term of any iterable."""
iterator = iter(iterable)
for _ in zip(range(n), iterator):
pass
return next(iterator)


Lastly, I would use existing tools to perform the iteration for me: itertools.islice is exactly what we need. In fact, your very function is already listed in the itertools recipes:

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


I would however remove the default value, especially in the case of an infinite iterable:

def nth(iterable, n):
"Returns the nth item or raise StopIteration"
return next(islice(iterable, n, None))


A completely unrelated note, but I would move the yield x at the beginning of the loop in triangle_stream so that you don't have to write it twice:

def triangle_stream():
""" Infinite stream of triangle numbers. """
x = step = 1
while True:
yield x
step += 1
x += step


And since step starts at 1, and never stop incrementing one by one, I would use itertools.count(1):

def triangle_stream():
""" Infinite stream of triangle numbers. """
x = 0
for step in itertools.count(1):
x += step
yield x


Now triangle_stream only accumulates values from itertools.count(1), well time to make use of itertools.accumulate:

def triangle_stream():
""" Infinite stream of triangle numbers. """
return itertools.accumulate(itertools.count(1))


or, to make it more obvious that you get a generator by calling triangle_stream:

def triangle_stream():
""" Infinite stream of triangle numbers. """
yield from itertools.accumulate(itertools.count(1))

• To further simplify: triangle_stream = itertools.accumulate(itertools.count(1)) Nov 3, 2016 at 14:12
• @200_success Correct me if I’m wrong, but this way you wouldn't be able to iterate several times over it, right? Nov 3, 2016 at 14:14

Which of these two nth_term functions is more efficient?

There are two types of efficacy, memory and speed. In terms of memory, they're about equal. In terms of speed, the one that doesn't use next is better. To test this you can use timeit:

>>> from timeit import timeit
>>> timeit('fn(gen, a)', 'from __main__ import nth_term_for as fn, triangle_stream as gen; a = 100', number=100000)
3.7175619061597445
>>> timeit('fn(gen, a)', 'from __main__ import nth_term_next as fn, triangle_stream as gen; a = 100', number=100000)
3.855789591803557


However, if i == n - 1: will reduce the speed of nth_term_for, and so if you change the code to remove that minus from the loop you can improve the performance. Take:

def nth_term_for_one(generator, n):
""" Returns the nth term of an infinite generator. """
for i, x in enumerate(generator(), 1):
if i == n:
return x


Finally, I wouldn't use any of the above functions. Instead I'd use the itertools recipe nth. Where I named it nth_term_iter for a comparison.

There's a problem with timeit, where it can range quite drastically even on the same input. And so I wrote a small program before to output graphs, that show the average, and the interquartile range of functions. This showed that the best is hands down itertools.islice. Here is the output that it generated:

• Is the graphing program open source? If so could you post a link to it? It looks like a useful script to have around. Nov 7, 2016 at 14:28
• @Caridorc No, ): I'll put it on GitHub later, and will ping you a link Nov 7, 2016 at 14:31
• @Caridorc Yo, I've put it on GitHub and made a question. MIT or CC-BY-SA. Nov 7, 2016 at 23:35
• I have taken a look at it, it looks very feature rich. I have also written a similar graphing program github.com/CAridorc/Plot-Benchmark but it is much more basic :) Nov 8, 2016 at 22:09