Delgan
Delgan

Reputation: 19697

Generator expression including conditional test without calling a function twice?

Suppose I have a function which performs some heavy calculations.

def f(x):
    ... 
    return result

Then I have a list with values I want to pass to f():

my_list = [2, 98, 4, 34, 23, 11]

I would like to find the first element x in this list which validates a condition on f(x) (let's say for example f(x) != 0), and get this computed result.

Basically, I would write a for loop like this:

def first_match(my_list):
    for x in my_list:
        r = f(x)
        if r != 0:
            return r

I would like to know if there is a way to get the same result using a generator expression?

What I thought so far was something like this:

r = next(f(x) if f(x) != 0 for x in my_list)

The issue is that this calls f() twice.

Upvotes: 1

Views: 671

Answers (5)

Cyphase
Cyphase

Reputation: 12022

There is a way to use generator expressions, as already mentioned by Blckknght, but there's also nothing wrong with the first_match() function in your question. Not everything has to be a one-liner :).

Here's another way to do it though:

try:
    from itertools import ifilter as filter  # Python 2
except ImportError:
    pass  # Python 3

predicate = lambda x: x != 0
r = next(filter(predicate, (f(x) for x in my_list)), None)

The second argument to next() (None in the above code) is what to return if there's nothing in the iterator that you pass in as the first argument. If you don't specify it and the iterator is exhausted, a StopIteration exception will be raised. You can change that to whatever you want.

Upvotes: 4

Terry Jan Reedy
Terry Jan Reedy

Reputation: 19219

Comprehensions are possibly nested map functions with option filtration of inputs. In this case, you want to the filter the output, so do so. Let y = f(x) = xx and let c(y) = (y > 100). Then (y for y in (xx for x in my_list) if y > 100) gives you the filtered outputs. As you noticed, stopping a filter on the first true element (a first-true search) can be done with next.

my_list = [2, 98, 4, 34, 23, 11]
ge = (y for y in (x*x for x in my_list) if y > 100)

print(next(ge))
# 9604

Upvotes: 1

zyxue
zyxue

Reputation: 8918

Maybe use itertools.ifilter

from itertools import ifilter
next(ifilter(lambda x: f(x) != 0, my_list))

Upvotes: 1

Blckknght
Blckknght

Reputation: 104842

You can use a nested generator expression to avoid the double function call:

next(y for y in (f(x) for x in my_list) if y != 0)

Upvotes: 7

msw
msw

Reputation: 43527

In general, generator expressions can be constructed as a generator function. For example, this f(x) has lazy evaluation as you'd expect from a generator:

def f(x):
    for i in x:
        yield i

Yes, this is a null delegation, but your example can be employed for your function (which I don't understand the intent of so didn't guess).

Upvotes: 0

Related Questions