Numpy：快速查找值的第一个索引

Question

我该怎么找到一个数字在Numpy数组中第一次出现的位置呢？速度对我来说很重要。我不想要那些会扫描整个数组的答案，因为它们在找到第一次出现时不会停止：

itemindex = numpy.where(array==item)[0][0]
nonzero(array == item)[0][0]

注意1：那个问题里的答案似乎都不相关 有没有Numpy函数可以返回数组中某个元素的第一个索引？

注意2：我更喜欢使用C语言编译的方法，而不是Python循环。

Answer 1

如果你想找到第一个非零的元素，可以使用以下方法：

idx = x.view(bool).argmax() // x.itemsize
idx = idx if x[idx] else -1

这个方法是一个非常快的“纯numpy”解决方案，但在一些情况下会出现问题，下面会讨论这些情况。

这个方法利用了一个事实：几乎所有数字类型表示的零都是由0字节组成的。这个道理同样适用于numpy的bool类型。在最近的numpy版本中，argmax()函数在处理bool类型时使用了短路逻辑。bool的大小是1字节。

所以你需要：

• 将数组视为bool类型。这样不会创建副本。
• 使用argmax()来找到第一个非零字节，这个过程会利用短路逻辑。
• 通过对偏移量进行整数除法（使用//运算符），将这个字节的偏移量转换为第一个非零元素的索引，除数是单个元素的字节大小（x.itemsize）。
• 检查x[idx]是否真的非零，以确定是否存在非零元素。

我对比了这个方法和numba的解决方案，并构建了np.nonzero。

import numpy as np
from numba import jit
from timeit import timeit

def find_first(x):
    idx = x.view(bool).argmax() // x.itemsize
    return idx if x[idx] else -1

@jit(nopython=True)
def find_first_numba(vec):
    """return the index of the first occurence of item in vec"""
    for i in range(len(vec)):
        if vec[i]:
            return i
    return -1


SIZE = 10_000_000
# First only
x = np.empty(SIZE)

find_first_numba(x[:10])

print('---- FIRST ----')
x[:] = 0
x[0] = 1
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0][0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=1000), 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=1000), 'ms')

print('---- LAST ----')
x[:] = 0
x[-1] = 1
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0][0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=100)*10, 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=100)*10, 'ms')

print('---- NONE ----')
x[:] = 0
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=100)*10, 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=100)*10, 'ms')

print('---- ALL ----')
x[:] = 1
print('ndarray.nonzero', timeit(lambda: x.nonzero()[0][0], number=100)*10, 'ms')
print('find_first', timeit(lambda: find_first(x), number=100)*10, 'ms')
print('find_first_numba', timeit(lambda: find_first_numba(x), number=100)*10, 'ms')

在我的机器上，结果是：

---- FIRST ----
ndarray.nonzero 57.63976670001284 ms
find_first 0.0010841979965334758 ms
find_first_numba 0.0002308919938514009 ms
---- LAST ----
ndarray.nonzero 58.96685277999495 ms
find_first 5.923203580023255 ms
find_first_numba 8.762269750004634 ms
---- NONE ----
ndarray.nonzero 25.13398071998381 ms
find_first 5.924289370013867 ms
find_first_numba 8.810063839919167 ms
---- ALL ----
ndarray.nonzero 55.181210660084616 ms
find_first 0.001246920000994578 ms
find_first_numba 0.00028766007744707167 ms

这个解决方案比numba快33%，而且是“纯numpy”的。

缺点是：

• 不适用于numpy接受的类型，比如object。
• 在float或double计算中，偶尔会出现负零，这种情况下会失败。

Answer 2

我做了几个方法的性能测试：

• argwhere
• nonzero，就像问题中提到的那样
• .tostring()，参考@Rob Reilink的回答
• Python循环
• Fortran循环

可以查看Python和Fortran的代码。我跳过了一些效果不佳的方法，比如转换成列表。

结果是以对数尺度展示的。X轴表示针的位置（如果针在数组的后面，找到它会花更长时间）；最后一个值是数组中没有的针。Y轴是找到它所需的时间。

这个数组有100万个元素，测试进行了100次。结果还是有点波动，但整体趋势很明显：Python和f2py在第一个元素时就停止了，所以它们的表现不同。如果针不在前1%的位置，Python会变得很慢，而f2py则很快（但需要编译）。

总结一下，f2py是最快的解决方案，特别是当针出现得比较早的时候。

虽然它不是内置的，这有点烦人，但其实只需要2分钟的工作。把这个添加到一个叫search.f90的文件中：

subroutine find_first(needle, haystack, haystack_length, index)
    implicit none
    integer, intent(in) :: needle
    integer, intent(in) :: haystack_length
    integer, intent(in), dimension(haystack_length) :: haystack
!f2py intent(inplace) haystack
    integer, intent(out) :: index
    integer :: k
    index = -1
    do k = 1, haystack_length
        if (haystack(k)==needle) then
            index = k - 1
            exit
        endif
    enddo
end

如果你要找的不是integer类型，只需更改类型。然后使用以下命令编译：

f2py -c -m search search.f90

之后你就可以在Python中这样做：

import search
print(search.find_first.__doc__)
a = search.find_first(your_int_needle, your_int_array)

Answer 3

虽然对你来说可能已经太晚了，但为了将来参考：使用 numba (1) 是最简单的方法，直到 numpy 自己实现这个功能。如果你使用的是 Anaconda 的 Python 版本，它应该已经安装好了。

这段代码会被编译，所以运行起来会很快。

@jit(nopython=True)
def find_first(item, vec):
    """return the index of the first occurence of item in vec"""
    for i in xrange(len(vec)):
        if item == vec[i]:
            return i
    return -1

然后：

>>> a = array([1,7,8,32])
>>> find_first(8,a)
2