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