List comprehension with condition

selurvedu picture selurvedu · Jun 27, 2014 · Viewed 18.5k times · Source

I have a simple list.

>>> a = [0, 1, 2]

I want to make a new list from it using a list comprehension.

>>> b = [x*2 for x in a]
>>> b
[0, 2, 4]

Pretty simple, but what if I want to operate only over nonzero elements?

'if' needs 'else' in list comprehensions, so I came up with this.

>>> b = [x*2 if x != 0 else None for x in a]
>>> b
[None, 2, 4]

But the desirable result is.

>>> b
[2, 4]

I can do that this way

>>> a = [0, 1, 2]
>>> def f(arg):
...     for x in arg:
...         if x != 0:
...             yield x*2
... 
>>> list(f(a))
[2, 4]

or using 'filter' and a lambda

>>> a = [0, 1, 2]
>>> list(filter(lambda x: x != 0, a))
[1, 2]

How do I get this result using a list comprehension?

Answer

Joran Beasley picture Joran Beasley · Jun 27, 2014
b = [x*2 for x in a if x != 0]

if you put your condition at the end you do not need an else (infact cannot have an else there)