如何根据其他列表的值过滤列表

Question

例如,我有两个相同长度的列表:

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]

该函数的作用是返回一个元素的列表,其列表b中相同索引的相应元素等于第三个参数.

在Matlab中,我将做一些像(b == 1),a(b == 2),b(b == 2)一样简单的事情.在Python中实现这一目标的最有效方法是什么？

Answer 1

如果你想要非常类似Matlab的功能,你可以使用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])

如果做不到这一点,我可能只是使用列表理解:

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

将它包装在一个函数中是微不足道的:

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

Answer 2

我会用itertools.compress()一个沿发电机表达要做到这一点:

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

请注意,这将返回一个生成器.如果您需要列表,请将其包装在list()通话中.

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

值得注意的是,这应该是好的和快速的 - itertools旨在成为一个快速的模块.