`elif` in list comprehension conditionals

Question

Can we use elif in list comprehension?

Example :

l = [1, 2, 3, 4, 5]
for values in l:
    if values==1:
        print 'yes'
    elif values==2:
        print 'no'
    else:
        print 'idle'

Can we include the elif in our list comprehension, in a similar fashion to the code above?

For example, an answer like:

['yes', 'no', 'idle', 'idle', 'idle']

Up until now, I have only used if and else in list comprehension.

Solution 1

Python's conditional expressions were designed exactly for this sort of use-case:

>>> l = [1, 2, 3, 4, 5]
>>> ['yes' if v == 1 else 'no' if v == 2 else 'idle' for v in l]
['yes', 'no', 'idle', 'idle', 'idle']

Solution 2

>>> d = {1: 'yes', 2: 'no'}
>>> [d.get(x, 'idle') for x in l]
['yes', 'no', 'idle', 'idle', 'idle']

Solution 3

You can, sort of.

Note that when you use sytax like:

['yes' if v == 1 else 'no' for v in l]

You are using the ternary form of the if/else operator (if you're familiar with languages like C, this is like the ?: construct: (v == 1 ? 'yes' : 'no')).

The ternary form of the if/else operator doesn't have an 'elif' built in, but you can simulate it in the 'else' condition:

['yes' if v == 1 else 'no' if v == 2 else 'idle' for v in l]

This is like saying:

for v in l:
    if v == 1 :
        print 'yes'
    else:
        if v == 2:
            print 'no'
        else:
            print 'idle'

So there's no direct 'elif' construct like you asked about, but it can be simulated with nested if/else statements.

Solution 4

Maybe you want this:

l = [1, 2, 3, 4, 5] 
print ([['idle','no','yes'][2*(n==1)+(n==2)] for n in l])

Solution 5

You can use list comprehension is you are going to create another list from original.

>>> l = [1, 2, 3, 4, 5]
>>> result_map = {1: 'yes', 2: 'no'}
>>> [result_map[x] if x in result_map else 'idle' for x in l]
['yes', 'no', 'idle', 'idle', 'idle']

View more solutions

Comments

[jdi]

I think this form is a lot easier to digest than trying to do a really long and complicated if/else logic within the list comp

[Raymond Hettinger]

@jdi Though conditional-expressions may not be to your taste, they were specifically designed to handle if-elif-elif-else chains, just as the OP requested. They aren't hard to learn and can gracefully handle situations that aren't as amenable to dictionary lookup logic: 'A' if grade>=90 else 'B' if grade>=80 else 'C' if grade>=70 else 'F'.

[Chris_Rands]

If there an advantage of defining d outside the comprehension?

[holdenweb]

There's some interesting history in the syntax. For many years before their introduction "tertiary expressions" were one of the five most-requested changes in the language. Since Guido van Rossum explicitly designed it as a statement-based language, he firmly resisted for a long time (tertiary expressions, and particularly their abuse, are sources of much obscurity in code). When he finally succumbed, he announced he had deliberately chosen a syntax that discouraged overuse. As usual, he did an elegant design job nevertheless.

[holdenweb]

Ternary, dammit (he wrote, noticing his dyslexic mistake too late to edit).

[Hai Vu]

While I up-vote this answer, I want to mention this: for 1 pair of if/else is easy to read, 2 pairs: it's getting harder to understand. Don't even mention 3 pairs. If the expression needs 3 or more pairs, a dictionary or a separate function will make things easier to read and understand.

[Alvaro Cavalcanti]

I would like to add not a solution for this problem, but a reminder of clean code: since this list comprehension has three conditionals, it could probably be refactored into a more descriptive method. My point is this: martinfowler.com/bliki/FunctionLength.html :)

[Tim Skov Jacobsen]

The reason I like the list comprehension better is that it reads just like English. Even a non-programmer would be able to understand what it does. With this solution you have to understand the dict.get() method.

[dTanMan]

I stumbled upon a case where I needed an elif, but only two values. Using this example, I would've needed just ['yes', 'no'] to be made. To do this, you can do: ['yes' if v == 1 else 'no' for v in l if values in [1,2]]. I currently can't think of a cleaner way to do this.

[mortysporty]

Great explanation... +1 for including the equivalent expression using the regular if else syntax.

[Christian Herenz]

This is especially useful if there are even more than three options.