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

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

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

Answer 1

函数numpy.in1d似乎做你想做的事。唯一的问题是它只适用于一维数组，所以你应该像这样使用它:

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)

我不知道为什么这仅限于一维数组。看着它的source code ，它首先似乎将两个数组展平，然后进行一些巧妙的排序技巧。但没有什么能阻止它在最后再次展开结果，就像我在这里必须手动做的那样。