查找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

函数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数组.看看它的源代码,它首先似乎压平了两个数组,之后它做了一些巧妙的排序技巧.但是没有什么可以阻止它再次在结束时取消结果,就像我必须在这里手工做的那样.

Answer 2

NumPy 0.13+

从 NumPy v0.13 开始，您可以使用np.isin，它适用于多维数组：

>>> 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 0.13 之前的版本

接受的答案np.in1d仅适用于一维数组，并且需要重新整形以获得所需的结果。这适用于 v0.13 之前的 NumPy 版本。

Answer 3

@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)

这个解决方案的con contained()是为每个元素调用的(即在python-space中),这使得它比纯粹的numpy解决方案慢得多.