Numpy中相当于list.index的函数

Question

在一个被多次调用的底层函数中，我需要使用numpy数组来实现类似于Python中list.index的功能。该函数需要在找到第一个值时返回该值，并在未找到目标值时抛出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中编写解决方案。

- lothario

---

你有一个numpy数组还是什么？ - Felix Kling

是的，我编辑了问题以澄清。 - lothario

你是否对代码进行了分析，以确保“where”是瓶颈所在。你可以展示你的代码的这一部分。据我所知，“numpy”中不存在你正在寻找的功能。谢谢。 - eat

我赞同@eat的观点。我认为numpy中没有一种方法可以做到你要求的事情，特别是返回ValueError。如果你想避免使用python循环，我建议你在cython中编写自己的函数，这应该会很快并且正好做你想要的事情。我也同意你应该分析你的代码，并查看使用nonzero或where再找到min索引是否真正成为了代码的瓶颈。相反，如果你多次调用该函数，问题可能在于你需要弄清楚是否可以使用numpy来避免多次调用，因为单个数组操作可能就足够了。 - JoshAdel

你拥有一个非常独特的情况。我从未在numpy中使用where或者在matlab中使用find遇到过任何严重的性能问题。(尽管仅仅使用一些普通的逻辑索引就足够解决问题) 除非你拟定一个非常具体的解决方案(适用于你的情况)，否则我不会期望从cython中获得任何主要改进。然而，请确保向我们展示您当前存在的瓶颈代码？谢谢。 - eat

Answer 1

请查看我在原帖中的评论以了解注意事项，但总的来说，我会执行以下操作：

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

- JoshAdel

Answer 2

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

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

如果不是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

- Damjan Banic

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]),)

您可能也会对此感兴趣：Where (http://www.scipy.org/Numpy_Example_List_With_Doc#where)。

- Joe

---

谢谢，但是 y >= 5 和 nonzero(y) 会遍历每个数组元素 -- 我正在寻找一个函数，它能在找到第一个索引时立即返回。 - lothario

这是正确的方向，如果你使用np.nonzero(y == index_value)[0][0]，它似乎可以在numpy数组中复制y.index(index_value)的功能，尽管我没有进行过广泛的测试。 - VectorVictor

Answer 4

你可以使用Cython编写代码，然后从Python脚本中导入。无需将整个项目迁移到Cython。

# paste into: indexing.pyx
def index(long[:] lst, long value):
    cdef int i
    for i in range(len(lst)):
        if lst[i] == value:
            return i
    raise ValueError

# import in your .py code
import pyximport
pyximport.install()
from indexing import index

# example
from numpy import zeros
a = zeros(10**6, int)
a[-1] = 1

index(a, 1)
Wall time: 6.07 ms
999999

index(a, 0)
Wall time: 38.1 µs
0

- ArekBulski

Answer 5

我遇到这个问题的唯一解决方法是将numpy数组转换为列表：

a = numpy.arange(3)
print(list(a).index(2))

>>> 2

- mmrocha

Answer 6

NumPy的searchsorted与列表的index非常相似，除了它需要一个排序好的数组并且表现更为数字化。最大的区别在于你不需要精确匹配，而且可以从左侧或右侧开始搜索。看一下以下示例，以便了解它的工作原理：

import numpy as np
a = np.array([10, 20, 30])

a.searchsorted(-99) == a.searchsorted(0) == a.searchsorted(10)
# returns index 0 for value 10

a.searchsorted(20.1) == a.searchsorted(29.9) == a.searchsorted(30)
# returns index 2 for value 30

a.searchsorted(30.1) == a.searchsorted(99) == a.searchsorted(np.nan)
# returns index 3 for undefined value

对于返回索引 3 的最后一种情况，您可以根据自己的需求进行处理。从函数的名称和意图中我推断它会在找到第一个合适的索引后停止。

- Mike T

---

我不知道它是否可以这样使用。它依赖于二分查找，而且原帖并没有指定数组已经排序。 - codewarrior

list.index 依赖于排序数组，文档中的“二分搜索”部分只是为那些混合数据类型的人提供警告（即 a = np.array([2.1, 3.1, 4.1], dtype=np.float16); a.searchsorted(np.float32(3.1)) 返回值为 2，而不是期望中的 1）。 - Mike T

list.index 不依赖于已排序的数组。 - codewarrior

没错，list.index返回第一个索引（如果有的话）。是的，这个解决方案只适用于已排序的数组。 - Mike T