Numpy中list.index的等价方法

Question

在一个被多次调用的低级函数中，我需要实现类似于Python中列表的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

我找到的最接近你所问的东西是 nonzero。听起来可能有点奇怪，但文档上看起来它可能会给你想要的结果。

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) 这个也可能对你有帮助。

Answer 2

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

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

如果不是一维的，你可以通过数组这样查找：

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

请查看我对提问者问题的评论，里面有一些注意事项，但一般来说，我会这样做：

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循环要快得多。如果坚持在找到第一个值时就停止搜索，可能并没有太大意义。