Multiple If/else in list comprehension in Python

Question

I have this

s = ['son','abc','pro','bro']
b = ['son','bro']
c = ['pro','quo']

The expected output is this. Where items in the output are index(item_in_s) if it is present in list b. Or index(item_in_s)+10 if an item is in c.

[0,12,3]

I tried this:

index_list = [s.index(item) if item in b else s.index(item)+10 if item in c for item in s]
print(index)

But apparently this is a syntax error. So I tried this:

index_list = [s.index(item) if item in b else s.index(item)+10 for item in s if item in c]
    print(index)

Output:

[12]

This just changes the whole logic. Although I could do this

fin = [s.index(item) if item in b else s.index(item)+10 if item in c  else '' for item in s]
fin = [item for item in fin if item!='']
print(fin)

Desired output obtained:

[0, 12, 3]

But how do I obtain what I want in list comprehension itself or is there something like else continue in list comprehensions?

Solution 1

Fundamentally, a list-comprehension forces you to be very inefficient:

>>> [i if item in b else i + 10 if item in c else None for i, item in enumerate(s) if item in b or item in c]
[0, 12, 3]

This has to check the membership item in b and c twice each in the worst-case if you want that output. Instead, just use a for-loop:

>>> index_list = []
>>> for i, item in enumerate(s):
...     if item in b:
...         index_list.append(i)
...     elif item in c:
...         index_list.append(i + 10)
...
>>> index_list
[0, 12, 3]
>>>

Simple, readable, straight-forward and Pythonic.

Solution 2

already good answers given, here is one which is not mentioned yet:

fin = [item for item in ([s.index(item) if item in b else s.index(item)+10 if item in c  else '' for item in s]) if item!='']
print(fin)

basically it is a combination of the original 2 lines of code:

fin = [s.index(item) if item in b else s.index(item)+10 if item in c  else '' for item in s]
fin = [item for item in fin if item!='']

Not necessarily "faster or better", just a combination which was not given before. When doing list comprehensions you always have the risk that you iterate more than you really need to.

A better solution would be a for loop, which is given in another answer.

Solution 3

index_list = [s.index(item) if item in b else s.index(item) + 10 for item in s if item in b or item in c]

We make sure it's in either b or c, and then the index will be either of the cases.

Solution 4

your solution glitch can be solved by avoiding the elements that doesn't exist in list b&c.

you can do this by creating new list and applying simple set operation

check this little change in your solution.

fin = [s.index(item) if item in b else s.index(item)+10 if item in c  else '' for item in list(set(b+c)&set(s))]

by doing this your conditinal statement else is never going to execute cause list on which you are iterating only has element that either list in b or c only.

View more solutions