CODEWITHSUNDEEP
python
running man
running man

Reputation: 1467

Python list comprehension for if else statemets

How to express by using list comprehension? Newbie needs help. Thanks a lot. code below:

lst = ['chen3gdu',2,['chengdu','suzhou']]
result = []
for elem in lst:
    if type(elem) == list:
        for num in elem:
            result.append(num)
    else:
        result.append(elem)

Upvotes: 6

Views: 641

Answers (2)

Dhruv Pal
Dhruv Pal

Reputation: 33

the answer from Martin Pieters(here) is great however i would like to suggest that in the final code sample of that answer which is 

def flatten(irregular_list):
    for elem in irregular_list:
        if isinstance(elem, list):
            yield from elem
        else:
            yield elem

if we tweak this to 

def flatten(irregular_list):
    for elem in irregular_list:
        if isinstance(elem, list):
            yield from flatten(elem)
        else:
            yield elem

then it can flatten to give the result of list containing only non-list elements even if the question list contain multiple list or "lists in a list" element.

PS: i was going to just comment this thing but i found out that i don’t have enough reputation.

Upvotes: 2

Martijn Pieters
Martijn Pieters

Reputation: 1123830

This isn't that suitable for a list comprehension, but you can achieve it by special-casing when you don't have a list, wrapping such elements in a list for iteration:

result = [num for elem in lst for num in ([elem] if not isinstance(elem, list) else elem)]

which, written out to the same for you were using, plus an extra variable to call out the conditional expression I used, is the equivalent of:

result = []
for elem in lst:
    _nested = [elem] if not isinstance(elem, list) else elem
    for num in _nested:
        result.append(num)

You might want to encapsulate flattening the irregular structure, in a generator function:

def flatten(irregular_list):
    for elem in irregular_list:
        if isinstance(elem, list):
            yield from elem
        else:
            yield elem

and then use that in list comprehensions and such, with additional operations. For just flattening, passing the generator function to list() is cleaner, e.g. result = list(flatten(lst)).

Upvotes: 9

Related Questions