查找NumPy数组与值列表中的任何值相等的位置

Question

我有一个整数数组，希望找到该数组与多个值列表中的任何值相等的位置。

这可以通过单独处理每个值，或者在一个循环中使用多个“或”语句来实现，但我觉得必须有更好/更快的方法来实现。我实际上是在处理大小为4000 x 2000的数组，但这里有一个问题的简化版：

fake = arange(9).reshape((3,3))

array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

want = (fake==0) + (fake==2) + (fake==6) + (fake==8)

print want 

array([[ True, False,  True],
       [False, False, False],
       [ True, False,  True]], dtype=bool)

我想要的是一种从单个命令获取want的方法，该命令涉及fake和值列表[0, 2, 6, 8]。

我假设有一个包已经包含了这个内容，这将比我用Python编写一个带循环的函数快得多。

Answer 1

@Bas的答案可能就是你要找的。但这里有另一种方法，使用numpy的vectorize技巧：

import numpy as np
S = set([0,2,6,8])

@np.vectorize
def contained(x):
    return x in S

contained(fake)
=> array([[ True, False,  True],
          [False, False, False],
          [ True, False,  True]], dtype=bool)

这个解决方案的优点是为每个元素（即在python空间中）调用contained()，这使得这个过程比纯numpy解决方案慢得多。

Answer 2

函数numpy.in1d似乎做了您想要的事情。唯一的问题是它只在1d数组上工作，所以您应该这样使用它：

In [9]: np.in1d(fake, [0,2,6,8]).reshape(fake.shape)
Out[9]: 
array([[ True, False,  True],
       [False, False, False],
       [ True, False,  True]], dtype=bool)

我不知道为什么这只限于1d数组。看看它的source code，它首先看起来是将两个数组展平，然后它做一些聪明的排序技巧。但没有什么能阻止它在最后再次破坏结果，就像我不得不在这里亲手做的那样。

Answer 3

纽比0.13+

从NumPy v0.13开始，您可以使用^{}，它可以在多维数组上工作：

>>> element = 2*np.arange(4).reshape((2, 2))
>>> element
array([[0, 2],
       [4, 6]])
>>> test_elements = [1, 2, 4, 8]
>>> mask = np.isin(element, test_elements)
>>> mask
array([[ False,  True],
       [ True,  False]])

NumPy pre-0.13

使用np.in1d接受的答案仅适用于1d数组，需要重新调整形状以获得所需的结果。这对v0.13之前的NumPy版本很好。