27

I was wondering, is it possible to put multiple if conditions in a list comprehension? I didn't find anything like this in the docs.

I want to be able to do something like this

ar=[]
for i in range(1,n):
  if i%4 == 0: ar.append('four')
  elif i%6 == 0: ar.append('six')
  else: ar.append(i)

using a list comprehension. How can I do it?

Is this even possible? If its not, what would be the most elegant (pythonic) way to accomplish this?

Rushil Paul
  • 2,010
  • 4
  • 26
  • 39

3 Answers3

47

How about

ar = [('four' if i % 4 == 0 else ('six' if i % 6 == 0 else i)) for i in range(1, n)]

For example, if n = 30 this is

[1, 2, 3, 'four', 5, 'six', 7, 'four', 9, 10, 11, 'four', 13, 14, 15, 'four', 17, 'six', 19, 'four', 21, 22, 23, 'four', 25, 26, 27, 'four', 29]

ETA: Here's how you could apply a list of conditions:

CONDITIONS = [(lambda i: i % 4 == 0, "four"), (lambda i: i % 6 == 0, "six"),
              (lambda i: i % 7 == 0, "seven")]

def apply_conditions(i):
    for condition, replacement in CONDITIONS:
        if condition(i):
            return replacement
    return i

ar = map(apply_conditions, range(0, n))
David Robinson
  • 77,383
  • 16
  • 167
  • 187
  • 2
    Thanks! I believe this way I can add any number of conditions to it. Though the code will be incomprehensible! – Rushil Paul Apr 22 '12 at 23:05
  • See my edit. I'm sure you can figure out the remainder (but if not I can help) – David Robinson Apr 22 '12 at 23:26
  • (My original neglected to make them lambda statements, current version would work though) – David Robinson Apr 22 '12 at 23:28
  • Ok so that's a `list` of conditions, but how do I use them in the list comprehension? And also, to put *i* in the list (when *i* is not divisible by any number), the code should be such that it knows when none of the conditions in `conditions` are true. – Rushil Paul Apr 23 '12 at 14:35
  • Yes this is good. I didn't know about a 2 variable for loop. Thanks! – Rushil Paul Apr 23 '12 at 15:48
  • 2
    That's just tuple unpacking- it could also have been done as `for c in conditions: if c[0](i): return c[1]` (the unpacked version is just clearer) – David Robinson Apr 23 '12 at 15:56
15

You can put you logic in a separate function, and then have the elegance of the list comprehension along with the readability of the function:

def cond(i):
    if i % 4 == 0: return 'four'
    elif i % 6 == 0: return 'six'
    
    return i

l=[cond(i) for i in range(1,n)]

If you have lots of conditions, it is usually easier to maintain a single dict rather than a big if/else ladder:

def cond(i):
    mkey={4:'four',6:'six'}
    return next((mkey[k] for k in mkey if i%k == 0), i)

This uses the default version of next to find if any integer key is a multiple of that key or the number itself, the default, if not.

Which could be a single comprehension if desired:

[next((v for k,v in {4:'four',6:'six'}.items() if i%k==0), i) for i  in range(1,10)]
dawg
  • 98,345
  • 23
  • 131
  • 206
5
ar = ["four" if i%4==0 else "six" if i%6==0  else i for i in range(1,30)]
RyanAbnavi
  • 358
  • 4
  • 6