CODEWITHSUNDEEP
pythonpython-2.7listlist-comprehension
scrap_metal
scrap_metal

Reputation: 399

Efficiently filtering out 'None' items in a list comprehension in which a function is called

I have a list comprehension that calls a function which potentially returns None.

f = lambda x: x if x < 3 else None
l = [f(x) for x in [1,2,3,4]]   # [1, 2, None, None]

I'd like to have the list comprehension as above, without the None entries.

What's a more efficient way to do the following, without creating the additional overhead, and while retaining the efficiency of the list comprehension?

filter(None, l)   # [1, 2]

Upvotes: 39

Views: 19662

Answers (2)

cottontail
cottontail

Reputation: 23081

It's already mentioned by Stephen but changing l in the OP into a generator expression (g below) allows you to evaluate f only once without creating an unnecessary list in the middle.

f = lambda x: x if x < 3 else None
l = [f(x) for x in [1,2,3,4]]            # <--- list       (inefficient)
g = (f(x) for x in [1,2,3,4])            # <--- generator  (efficient)
output = [x for x in g if x is not None]

You can use filter() on the generator too (for a solution closer to what OP had in mind):

output = list(filter(None, g))

N.B. filter(None, g) will filter out all falsy values, so it will filter out 0; if you don't want to do that, then probably the list comprehension above is the cleaner solution.

Python>=3.9

Since Python 3.9, creating a singleton list and looping over it has been optimized. So instead of using a generator expression (which may be exhausted by accident) or the walrus operator (which leaks the loop variable to the outer scope), you can create a nested loop where the second loop is over a singleton list:

output = [y for x in [1,2,3,4] for y in [f(x)] if y is not None]

Simple timeit test (on Python 3.12) shows that this is actually slightly faster than the other options (including the one using the walrus operator).

Upvotes: 1

Stephen Rauch
Stephen Rauch

Reputation: 49784

Add an if into your comprehension like:

l = [y for y in (f(x) for x in [1,2,3,4]) if y is not None]

By placing a Generator Expression inside the list comprehension you will only need to evaluate the function once. In addition the generator expression is a generator so takes no extra intermediate storage.

Python 3.8+

As of Python 3.8 you can use an Assignment Expression (:=) (AKA: Named Expressions or the Walrus Operator) to avoid multiple evaluations of f() like:

l = [y for x in [1,2,3,4] if (y := f(x)) is not None]

Upvotes: 83

Related Questions