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

【问题标题】：Find where a NumPy array is equal to any value in a list of values查找 NumPy 数组与值列表中的任何值相等的位置
【发布时间】：2013-10-23 18:30:59
【问题描述】：

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

这可以通过单独处理每个值或在循环中使用多个“或”语句来轻松完成，但我觉得必须有更好/更快的方法来做到这一点。我实际上正在处理大小为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 中编写一个带有循环的函数要快得多。

【问题讨论】：

标签： python arrays numpy

【解决方案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，它首先似乎将两个数组展平，然后它做了一些巧妙的排序技巧。但没有什么能阻止它在最后再次弄平结果，就像我必须在这里手工做的那样。

【讨论】：

嗯。我编写了这个非常简单的函数来完成这项工作：def EqualsAny(ar,vals): out=zeros(ar.shape,dtype=bool) for val in vals: out+=(ar==val) return out 我认为numpy.in1d 会更快，但实际上需要更长的时间（对于相同的结果）：In [11]: %timeit EqualsAny(badlabels,smallnum) 1 loops, best of 3: 519 ms per loop In [7]: %timeit in1d(badlabels, smallnum).reshape(badlabels.shape) 1 loops, best of 3: 871 ms per loop numpy.in1d 不应该更快，因为它是写的在 C?我没有正确使用%timeit 吗？
不，in1d 不是用 c 编写的，而是用 python 编写的，请参阅我提供的源代码的链接。它使用各种 numpy 函数，如 sort，希望用 C 编写。它甚至有一些优化算法，用于当 vals 很短时，这与您的解决方案非常相似（但使用 |= 而不是 @987654333 @)。我不知道为什么你的版本更快，这可能取决于两个输入的长度。

【解决方案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 pre-0.13

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

【讨论】：

【解决方案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)

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

【讨论】：