CODEWITHSUNDEEP
python
Ruofeng
Ruofeng

Reputation: 2350

How to filter a list based on another list's values

For example, I have two lists of the same length:

a = [[1,2],[2,2],[3,3],[4,2],[5,6]]
b = [1,2,2,3,1]

I want to have a function such that 

func(a,b,1) = [[1,2],[5,6]]
func(a,b,2) = [[2,2],[3,3]]
func(b,b,2) = [2,2]

What the function does is return a list of a's elements, whose corresponding elements of the same index in list b equal to the third argument.

In Matlab I will do something as easy as a(b==1), a(b==2), b(b==2). What is the most efficient way to achieve this in Python?

Upvotes: 6

Views: 11413

Answers (4)

Ashwini Chaudhary
Ashwini Chaudhary

Reputation: 251136

def func(a,b,n):
   return [a[i] for i,x in enumerate(b) if x==n]

a = [[1,2],[2,2],[3,3],[4,2],[5,6]]
b = [1,2,2,3,1]
print(func(a,b,1)) #[[1, 2], [5, 6]]
print(func(a,b,2)) #[[2, 2], [3, 3]]
print(func(b,b,2)) #[2, 2]

Upvotes: 1

Gareth Latty
Gareth Latty

Reputation: 89097

I would use itertools.compress() along with a generator expression to do this:

def func(a, b, c):
    return itertools.compress(a, (x == c for x in b))

Note that this will return a generator. If you need a list, wrap it in a list() call.

>>> import itertools
>>> def func(a, b, c):
...     return list(itertools.compress(a, (x == c for x in b)))
... 
>>> a = [[1,2],[2,2],[3,3],[4,2],[5,6]]
>>> b = [1,2,2,3,1]
>>> func(a, b, 1)
[[1, 2], [5, 6]]
>>> func(a, b, 2)
[[2, 2], [3, 3]]
>>> func(b, b, 2)
[2, 2]

It's also worth noting this should be nice and fast - itertools is designed to be a fast module.

Upvotes: 6

Junuxx
Junuxx

Reputation: 14281

This first finds the indices where b equals the target values ([i for i, x in enumerate(b) if x==c]), then uses those indices to get the elements you want out of a:

def func(a,b,c):
    return [a[j] for j in [i for i, x in enumerate(b) if x==c]]

>>> func(a,b,1)
[[1, 2], [5, 6]]
>>> func(a,b,2)
[[2, 2], [3, 3]]
>>> func(b,b,2)
[2, 2]

Upvotes: 4

DSM
DSM

Reputation: 353479

If you want very Matlab-like functionality, you could use numpy:

>>> import numpy
>>> a = [[1,2],[2,2],[3,3],[4,2],[5,6]]
>>> b = [1,2,2,3,1]
>>> a = numpy.array(a)
>>> b = numpy.array(b)
>>> a[b==1]
array([[1, 2],
       [5, 6]])
>>> a[b==2]
array([[2, 2],
       [3, 3]])
>>> b[b==2]
array([2, 2])

Failing that, I'd probably simply use a list comprehension:

>>> [i for i,j in zip(a,b) if j == 1]
[[1, 2], [5, 6]]
>>> [i for i,j in zip(a,b) if j == 2]
[[2, 2], [3, 3]]
>>> [i for i,j in zip(b,b) if j == 2]
[2, 2]

It'd be trivial to wrap this in a function:

>>> def func(a,b,x):
...     return [i for i,j in zip(a,b) if j == x]
... 
>>> func(a,b,2)
[[2, 2], [3, 3]]
>>>

Upvotes: 11

Related Questions