Numpy搜索子集中数组的元素

时间:2014-11-30 19:35:29

标签: python numpy

假设我有numpy数组

a = np.array([1,3,5,7,9,11,13])
b = np.array([3,5,7,11,13])

我想创建一个大小为a的布尔数组,其中每个条目都是True或False,具体取决于a的元素是否也在b中。

所以在这种情况下,我想要

a_b = np.array([False,True,True,True,False,True,True]). 

当b由一个元素组成== b [0]时,我可以这样做。当b的长度大于1时,有没有快速的方法。

1 个答案:

答案 0 :(得分:5)

使用numpy.in1d

In [672]: np.in1d([1,2,3,4], [1,2])
Out[672]: array([ True,  True, False, False], dtype=bool)

对于您的数据:

In [674]: np.in1d(a, b)
Out[674]: array([False,  True,  True,  True, False,  True,  True], dtype=bool)

根据文档,这在1.4.0或更高版本中可用。文档还描述了纯Python中操作的外观:

  对于1-D序列,

in1d可以被视为python关键字in的元素功能版本。 in1d(a, b)大致相当于np.array([item in b for item in a])

这个函数的文档值得阅读,因为有invert关键字参数和assume_unique关键字参数 - 在某些情况下,每个参数都非常有用。

我还发现使用np.vectorizeoperator.contains创建自己的版本很有意思:

from operator import contains
v_in = np.vectorize(lambda x,y: contains(y, x), excluded={1,})

然后:

In [696]: v_in([1,2,3, 2], [1, 2])
Out[696]: array([ True,  True, False,  True], dtype=bool)

因为operator.contains翻转了参数,我需要lambda来使调用约定与您的用例匹配 - 但如果可以先使用b调用,则可以跳过这个然后a

但请注意,您需要对excluded使用vectorize选项,因为您希望任何参数代表b序列(检查其中的成员资格的序列)实际保留为序列(因此,如果您选择不使用contains翻转lambda参数,那么您可能希望排除索引0而不是1。)

使用in1d的方式肯定会更快,并且是一种更好的方式,因为它依赖于众所周知的内置功能。但有时候知道如何使用operatorvectorize这些技巧很好。

您甚至可以为此创建Python Infix recipe instance,然后使用v_in作为“中缀”操作:

v_in = Infix(np.vectorize(lambda x,y: contains(y, x), excluded={1,}))
# even easier: v_in = Infix(np.in1d)

和示例用法:

In [702]: v_in([1, 2, 3, 2], [1, 2])
Out[702]: array([ True,  True, False,  True], dtype=bool)

In [704]: [1, 2, 3, 2] <<v_in>> [1, 2]
Out[704]: array([ True,  True, False,  True], dtype=bool)