Question

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?

Answer 1

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]]
>>>