Search code examples
pythoncountpython-itertoolsinfinite

Python summing itertools.count?

I'm having some trouble with the itertools.count function, and I don't quite understand what it does. I expect the code below to accomplish Project Euler problem 2.

I know that I could write this with a simple while loop, but is there a way to do it with a list comprehension? This code just freezes as I guess it's going to go infinity with count(). I would have hoped it would stop after x > MAX, but I know that won't happen. Is there a way to stop count in a generator expression like below? 

def fib(n):
    if (n <= 1): return 1
    else: return fib(n-1) + fib(n-2)


MAX = 4000000

infiniteFib = (fib(x) for x in count())

s = (x for x in infiniteFib if x < MAX and x % 2 == 0)

print sum(s)
Solution

  • We just need to tell the infiniteFib generator when to stop yielding elements. itertools provides a number of useful methods to help with this:

    less_than_max = itertools.takewhile(lambda x: x<MAX, infiniteFib))
even = itertools.ifilter(lambda x: x%2==0, less_than_max)
print sum(even)

    We get a generator for all the numbers yielded by infiniteFib, until one returned is greater than MAX. Then we filter that generator, choosing only the even numbers. And finally we can sum the result.