CODEWITHSUNDEEP
pythonpython-3.x
user8023293
user8023293

Reputation: 67

Filter list of lists based on another list of lists

I'm trying to filter a list of lists based on the values of another list of lists in Python.

For example, I have:

list1 = [[0,1,2],[0,2,3]]
list2 = [['a','b','c'],['a','b','b']]

and want to filter list1 so that it only contains values at the same indexes as 'a' in list2. So my desired result is

filteredList_a = [[0],[0]]

Similarly, filtering list1 to have only values at the same indexes as 'b' in list2 would give

filteredList_b = [[1],[2,3]]

I know how to do this for a single list,

>>> list1 = [0,1,2]
>>> list2 = ['a','a','b']
>>> [i for index, i in enumerate(list1) if list2[index] == 'a']
[0,1]

Upvotes: 4

Views: 377

Answers (1)

ggorlen
ggorlen

Reputation: 56865

Here's an extension of your list comprehension approach with nested list comprehensions and zip to avoid indexes:

def filter_by(a, b, target):
    return [[i for i, j in zip(x, y) if j == target] for x, y in zip(a, b)]


list1 = [[0,1,2],[0,2,3]]
list2 = [['a','b','c'],['a','b','b']]
print(filter_by(list1, list2, 'a'))
print(filter_by(list1, list2, 'b'))

Output:

[[0], [0]]
[[1], [2, 3]]

Try it!

Upvotes: 4

Related Questions