remove elements in one list according to the conditions on corresponding elements of another list

Question

I have two lists that have the same length, lstA, lstB, I want to remove elements that are not good based on some conditions, but at the meantime, I also want to remove the elements in lstB that have the same position with the unwanted one in lstA. like below:

lstA = ['good','good2', '']
lstB = [1,2,3]

I want '' in lstA and 3 in lstB removed:

results:
lstA = ['good','good2']
lstB = [1,2]

Currently, I am doing this:

lstB = [lstB[i] for i in range(len(lstB)) if lstA[i] != '']
lstA = [st for st in lstA if st != '']

Want a more elegant and pythonic way of doing this

Answer 1

Rather than two comprehensions, it'd be simpler to just use one for loop. Make new_list_a and new_list_b start empty, and if the item is good add to both at the same time. This is simpler to understand, as you see clearly that both are being filtered by the same thing.

new_list_a = []
new_list_b = []
for a, b in zip(list_a, list_b):
    if a != '':
        new_list_a.append(a)
        new_list_b.append(b)

However, it looks like you'd be better served if you only create one list. This means you can have a tuple of (a, b). And so could implement something like:

[(a, b) for a, b in zip(list_a, list_b) if a != '']

Answer 2

Maybe this? I should say that this works in python 3.6.

The following code creates a list of tuples from the elements of 'lstA' and 'lstB' excluding the tuple that would contain the offending element. The zip(*[]) unpacks the list and then zips the resulting tuples thus creating the desired groups. Please see this example:

l = [('a', 1), ('b', 2)]

# the star operator is doing the job of the following two lines
# it could be: l1, l2 = [*l]
l0 = l[0]
l1 = l[1]

print(list(zip(l0,l1)))

Then the list() converts the tuples to lists.

Here is the code:

lstA = ['good','good2', '']
lstB = [1,2,3]

lstA, lstB = list(zip(*[(st, lstB[idx]) for idx, st in enumerate(lstA) if st != '']))

Or you can just change the line that creates 'lstB' in your code. The difference here is the use of enumerate() which gives as the index of each element of 'lstA'. As the elements that you want to keep have the same index in both lists, we keep the elements of 'lstB' that correspond to the elements of 'lstA' and are not the undesirable empty string ''.

lstA = ['good','good2', '']
lstB = [1,2,3]

lstB = [lstB[idx] for idx, st in enumerate(lstA) if st != '']
lstA = [st for st in lstA if st != '']