Convert double for loop with break to list comprehension

Question

I have a type something like this:

class T:
    id: int
    data: Any

I have a list of Ts with unique ids. I am provided another list of Ts with new data. So I want a new list that replaces any of the items where I have a new T with that new T.

current = [T(1), T(2), T(3)]
new = [T(2)*]

output = [T(1), T(2)*, T(3)]

I have a working double for loop with a break that I’m trying to turn into a list comprehension if it’s cleaner.

output = []

for item in current_items:
    for new_item in new_items:
        if new_item.id == item.id:
            item = new_item
            break
    output.append(item)

This is what I tried but without the ability to break the inner loop after performing the if condition it obviously doesn’t work.

output = [
    new_item if new_item.id == item.id else item
    for item in current_items            
    for new_item in new_items
]

Asked By: JonFitt

||

Source

Answer 1

Basically what you are doing is setting the output to be new_items that are in the current items. I would suggest creating a set of current item ids and then just filtering new items that are in that set.

current_item_ids = {
    item.id
    for item in current_items
}

output = [
    item
    for item in new_items
    if item in current_item_ids
]

Answered By: Igor Dragushhak

Answer 2

Let’s feed next a generator expression finding elements in b with the same id as the current element in a, and default to the current element in a if no matching id is found in b.

>>> from dataclasses import dataclass
>>> @dataclass
... class T(object):
...   id: int
...   name: str
... 
>>> a = [T(1, "foo"), T(2, "bar"), T(3, "baz")]
>>> b = [T(2, "wooble")]
>>> [next((y for y in b if y.id == x.id), x) for x in a] 
[T(id=1, name='foo'), T(id=2, name='wooble'), T(id=3, name='baz')]

Using next will replicate the behavior of your loop with a break on finding a match. The entire b list will only be iterated if a match isn’t found.

If we didn’t care about efficiency, we could generate a list of all matching items in b and then take the first one if it’s not empty.

[c[0] if (c := [y for y in b if y.id == x.id]) else x for x in a]

But that both looks uglier and is potentially less efficient both in terms of runtime complexity and space, as it generates a bunch of useless lists.

Answered By: Chris

Convert double for loop with break to list comprehension

Question:

Answers: