Ryan Endacott
Ryan Endacott

Reputation: 9172

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)

Upvotes: 2

Views: 2149

Answers (5)

DSM
DSM

Reputation: 353119

You could use takewhile:

>>> from itertools import count, takewhile, imap
>>> sum(x for x in takewhile(lambda x: x < 4000000, imap(fib, count())) if x % 2 == 0)
4613732

Upvotes: 8

NPE
NPE

Reputation: 500377

How about:

def fib():
    a, b = 1, 1
    while True:
        yield b
        a, b = b, a+b

sum(f for f in itertools.takewhile(functools.partial(operator.ge, 4000000), fib()) if f % 2 == 0)

Or, pushing the parity check into the generator:

def even_fib():
    a, b = 1, 1
    while True:
        if b % 2 == 0: yield b
        a, b = b, a+b

sum(itertools.takewhile(functools.partial(operator.ge, 4000000), even_fib()))

Upvotes: 3

ceyko
ceyko

Reputation: 4852

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.

Upvotes: 5

Marius
Marius

Reputation: 60080

Here's another solution using takewhile, but non-recursively. Since the recursive solution requires calculating all the fibs less than n for each n, it's horrible slow.

def fib_gen(only_even=False):
    one = 1
    if not only_even:
        yield one
    two = 1
    if not only_even:
        yield two
    while True:
        next = one + two
        one = two
        two = next
        if only_even:
            if next % 2 == 0:
                yield next
        else:
            yield next

list(itertools.takewhile(lambda x: x < 4000000, fib_gen()))

Upvotes: 0

jimhark
jimhark

Reputation: 5046

Yeah, count() just keeps going, which isn't what you want. List comprehensions / iterator expressions don't have flexible exit conditions (but see @DSM's solution using takewhile).

I prefer just using while.

Here's my old answer to Euler 2:

def SumEvenFibonacci(limit):
        x = y = 1
        sum = 0
        while (sum <= limit):
                sum += (x + y)
                x, y = x + 2 * y, 2 * x + 3 * y
        return sum

ce = SumEvenFibonacci(4000000)
print ce

Upvotes: 0

Related Questions