python - list.index 的 Numpy 等价物

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 循环搜索整个数组时，它相对较快。坚持在找到第一个值时停止搜索可能在功能上无关紧要。