Numpy等效于list.index

Question

在一个被多次调用的低级函数中,我需要做相当于python的list.index,但是需要一个numpy数组.函数需要在找到第一个值时返回,否则引发ValueError.就像是:

>>> a = np.array([1, 2, 3])
>>> np_index(a, 1)
0
>>> np_index(a, 10)
Traceback (most recent call last):    
  File "<stdin>", line 1, in <module>
ValueError: 10 not in array

我希望尽可能避免使用Python循环. np.where不是一个选项,因为它总是遍历整个数组; 一旦找到第一个索引,我需要停止的东西.

---

编辑:与问题相关的一些更具体的信息.

• 大约90%的时间,我正在搜索的索引是阵列的前1/4到1/2.因此,这里可能存在2-4倍的加速因素.另外10%的时间值根本不在数组中.

• 我已经分析了一些东西,并且调用np.where是瓶颈,至少占总运行时间的50%.

• 提出一个并不重要ValueError; 它只需要返回一些显然表明该值不在数组中的东西.

我建议可能会在Cython中编写一个解决方案.

Answer 1

请参阅我对OP的问题的评论,但一般来说,我会做以下事情:

import numpy as np
a = np.array([1, 2, 3])
np.min(np.nonzero(a == 2)[0])

如果您要查找的值不在数组中,您将得到以下内容ValueError:

ValueError: zero-size array to ufunc.reduce without identity

因为你试图获取空数组的最小值.

我会分析这段代码并查看它是否是一个真正的瓶颈,因为通常当numpy使用内置函数而不是显式python循环搜索整个数组时,它相对较快.当发现第一个值时,坚持停止搜索可能在功能上无关紧要.

Answer 2

如果你的 numpy 数组是一维数组，也许可以尝试这样：

a = np.array([1, 2, 3])
print a.tolist().index(2)
>>> 1

如果不是 1d，您可以搜索数组，如下所示：

a = np.array([[1, 2, 3],[2,5,6],[0,0,2]])
print a[0,:].tolist().index(2)
>>> 1

print a[1,:].tolist().index(2)
>>> 0

print a[2,:].tolist().index(2)
>>> 2

Answer 3

我能找到的最接近您要求的东西是非零。这可能听起来很奇怪，但文档让它看起来可能有预期的结果。

http://www.scipy.org/Numpy_Example_List_With_Doc#nonzero

具体这部分：

a.nonzero()

返回非零元素的索引。

请参阅numpy.nonzero完整文档。

也可以看看

numpy.nonzero ：等效函数

>>> from numpy import *
>>> y = array([1,3,5,7])
>>> indices = (y >= 5).nonzero()
>>> y[indices]
array([5, 7])
>>> nonzero(y)                                # function also exists
(array([0, 1, 2, 3]),)

其中 (http://www.scipy.org/Numpy_Example_List_With_Doc#where) 您可能也感兴趣。