pythonlist-comprehension

Use list comprehension with if, else, and for-loop while only keeping list items that meet a condition


I use list comprehension to load only the images from a folder that meet a certain condition. In the same operation, I would also like to keep track of those that do not meet the condition. This is where I am having trouble.

If the if and else conditions are at the beginning, each iteration yields a resulting element. The items that are caught by else are replaced by None rather than being excluded from the result list.

I can't figure out how to add an else condition such that an operation can be done on the items caught by else without including them in the resulting list.

here is a simplified and generalized version of the code:

exclude_imgs = [1, 3]
final = [
    n
    for ix, n in enumerate(sorted(["img1", "img4", "img3", "img5", "img2"]))
    if ix + 1 not in exclude_imgs
]
[ins] In [4]: final
Out[4]: ['img2', 'img4', 'img5']

adding else condition to store excluded images:

exclude_imgs = [1, 3]
excluded = []
final = [
    n if ix + 1 not in exclude_imgs else excluded.append(n)
    for ix, n in enumerate(sorted(["img1", "img4", "img3", "img5", "img2"]))
]
[ins] In [6]: final
Out[6]: [None, 'img2', None, 'img4', 'img5']

[ins] In [7]: excluded
Out[7]: ['img1', 'img3']

How can I write this so that the results are as follows:

final: ['img2', 'img4', 'img5']
excluded: ['img1', 'img3']

?


Solution

  • Consider not using a list comprehension at all, precisely because you want to create two lists, not just one.

    exclude_imgs = [1, 3]
    excluded = []
    final = []
    
    for ix, n in enumerate(sorted(["img1", "img4", "img3", "img5", "img2"]), start=1):
        (excluded if ix in excluded_imgs else final).append(n)