python - how to replace if statements in a for loop (for better looking code)

Question

I'd like to remove repetitive code but I'm stuck. I want the for loop to iterate through a list and append the object to another list depending on if the list has "red","blue",or "green" in it. I'm also trying to remove None data and strings of the names of the list. Not quite sure what to do here but I'd really appreciate some help!

for item in sublist:
    if "red" in sublist:
        red_list.append(item)
    if "blue" in sublist:
        blue_list.append(item)
    if "green" in sublist:
        green_list.append(item)

red_list = list(filter(None, red_list)
blue_list = list(filter(None, blue_list)
green_list = list(filter(None, green_list)

red_list.remove("red")
blue_list.remove("blue")
green_list.remove("green")

I've tried appending the items to a dictionary instead of three separate lists and that has been helpful, but I'm still not sure what to do about the for loop and if statements.

What are you trying to achieve? The code repeats the same queries `if "red" in sublist` in every loop cycle, `independently on what item is. — Maciej Skorski, Apr 06 '23 at 07:22
Where the code says e.g. `if "red" in sublist:`, did you mean `if "red" in item:`? Otherwise, what is the intended purpose of checking the entire list, repeatedly for each element of the list? What actually is the overall goal of the code? — Karl Knechtel, Apr 06 '23 at 08:10

S3DEV · Accepted Answer · 2023-04-06T07:44:55.843

A defaultdict can be very helpful here. The defaultdict enables the creation of the specified dict-key-object (list) if it does not already exist, or uses the existing object's methods if it does exist.

On first iteration, the dict is empty, but the list.append method can be called on the 'empty' key, as the default object of the dict is a list. This SO post goes into greater detail.

For example:

from collections import defaultdict

sublist = ['green', 'red', 'red', 'red', 'blue', 'blue', 'green']

d = defaultdict(list)
for i in sublist:
    d[i].append(i)

Output:

defaultdict(list,
            {'green': ['green', 'green'],
             'red': ['red', 'red', 'red'],
             'blue': ['blue', 'blue']})

Which can be accessed just like a builtin dict:

>>> d['green']
['green', 'green']

Or, explicitly converted into a builtin dict:

>>> dict(d)
{'green': ['green', 'green'],
 'red': ['red', 'red', 'red'],
 'blue': ['blue', 'blue']}

Removing unwanted items

If the source list contains unwanted items, they can be removed as shown below.

The complete code:

from collections import defaultdict

sublist = ['green', 'red', 'red', 'red', 'blue', 'blue', 'green',
           None, 'Sam', 'Bob', 'Sue', None, None]

# Append items to dict key/items.
d = defaultdict(list)
for i in sublist:
    d[i].append(i)

# Define a list of items to be kept.
keep = ['blue', 'red', 'green']

# Iterate a *copy* of the dict and pop the unwanted items.
for k in d.copy():
    if k not in keep:
        d.pop(k)

Output:

defaultdict(list,
            {'green': ['green', 'green'],
             'red': ['red', 'red', 'red'],
             'blue': ['blue', 'blue']})

score 0 · Answer 2 · answered Apr 06 '23 at 07:23

0

You could do something like this? (untested)

colors = ['red', 'blue', 'green']
lists = {color: [] for color in colors}
    
for color, list_ in lists.items():
    for item in sublist:
        if color in sublist:
            list_.append(item)
    lists[color] = list(filter(None, list_).remove(color)

answered Apr 06 '23 at 07:23

Julien

13,986
5
29
53

Without an actual [mre] provided by the user to be able to test, I believe this code is equivalent to OP's but more concise and pythonic. If not please explain... – Julien Apr 06 '23 at 12:40

python - how to replace if statements in a for loop (for better looking code)

2 Answers2

Removing unwanted items