Numpy: 快速找到值的第一个索引

如何在 Numpy 数组中找到数字的第一个匹配项的索引？ 速度对我很重要。我对下面的答案不感兴趣，因为它们会扫描整个数组，并且在发现第一个匹配项时不会停止:

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

注1: 这个问题的答案似乎都不相关

Note 2: using a C-compiled method is preferred to a Python loop.

103481 次浏览
小开

我认为您遇到了一个问题，在这个问题中，不同的方法和一些关于数组的 先验的知识会非常有帮助。这种情况下，你有 X 的概率找到你的答案在第一个 Y 百分比的数据。把问题分解成希望获得好运，然后在 python 中使用嵌套的列表内涵或者其他东西。

使用 类型编写一个 C 函数来完成这个强力操作也不是太难。

我拼凑的 C 代码(index.c) :

long index(long val, long *data, long length){
long ans, i;
for(i=0;i<length;i++){
if (data[i] == val)
return(i);
}
return(-999);
}

还有蟒蛇:

# to compile (mac)
# gcc -shared index.c -o index.dylib
import ctypes
lib = ctypes.CDLL('index.dylib')
lib.index.restype = ctypes.c_long
lib.index.argtypes = (ctypes.c_long, ctypes.POINTER(ctypes.c_long), ctypes.c_long)


import numpy as np
np.random.seed(8675309)
a = np.random.random_integers(0, 100, 10000)
print lib.index(57, a.ctypes.data_as(ctypes.POINTER(ctypes.c_long)), len(a))

我得到92分。

将 Python 包装成一个适当的函数，就可以了。

对于这个种子，C 版本要快得多(大约20倍)(警告: 我不擅长计时)

import timeit
t = timeit.Timer('np.where(a==57)[0][0]', 'import numpy as np; np.random.seed(1); a = np.random.random_integers(0, 1000000, 10000000)')
t.timeit(100)/100
# 0.09761879920959472
t2 = timeit.Timer('lib.index(57, a.ctypes.data_as(ctypes.POINTER(ctypes.c_long)), len(a))', 'import numpy as np; np.random.seed(1); a = np.random.random_integers(0, 1000000, 10000000); import ctypes; lib = ctypes.CDLL("index.dylib"); lib.index.restype = ctypes.c_long; lib.index.argtypes = (ctypes.c_long, ctypes.POINTER(ctypes.c_long), ctypes.c_long) ')
t2.timeit(100)/100
# 0.005288000106811523
小开

As far as I know only np.any and np.all on boolean arrays are short-circuited.

在您的示例中，numpy 必须遍历整个数组两次，一次是为了创建布尔条件，第二次是为了查找索引。

在这种情况下，我的建议是使用 cython。我认为为这种情况调整示例应该很容易，特别是如果您不需要为不同的 dtype 和形状提供很大的灵活性。

小开
最佳答案

在 Numpy 2.0.0: https://github.com/numpy/numpy/issues/2269中安排了一个这方面的特性请求

小开

我的工作需要这个，所以我自学了 Python 和 Numpy 的 C 接口，并编写了自己的接口。http://pastebin.com/GtcXuLyd它只适用于1-D 数组，但适用于大多数数据类型(int、 float 或字符串) ，测试表明它再次比纯 Python-numpy 中的预期方法快20倍左右。

小开

您可以将数组转换为 list并使用它的 index()方法:

i = list(array).index(item)

据我所知，这是一个 C 编译的方法。

小开

如果列表为 解决了，则可以使用“对分”包实现对索引的 很快搜索。 是 O (log (n))而不是 O (n)。

bisect.bisect(a, x)

在数组 a 中找到 x，在排序的情况下肯定比任何 C 例程通过所有第一个元素(对于足够长的列表)更快。

有时候知道真好。

小开

您可以使用 array.tostring()并使用 find ()方法将布尔数组转换为 Python 字符串:

(array==item).tostring().find('\x01')

This does involve copying the data, though, since Python strings need to be immutable. An advantage is that you can also search for e.g. a rising edge by finding \x00\x01

小开

在排序数组的情况下，np.searchsorted工作。

小开

请注意，如果您正在执行一系列搜索，那么如果搜索维度不够大，那么从转换为字符串等聪明操作中获得的性能收益可能会在外部循环中丢失。查看使用上面提出的字符串转换技巧迭代 find1和沿着内轴使用 argmax 的 find2的性能(加上一个调整以确保非匹配返回为 -1)

import numpy,time
def find1(arr,value):
return (arr==value).tostring().find('\x01')


def find2(arr,value): #find value over inner most axis, and return array of indices to the match
b = arr==value
return b.argmax(axis=-1) - ~(b.any())




for size in [(1,100000000),(10000,10000),(1000000,100),(10000000,10)]:
print(size)
values = numpy.random.choice([0,0,0,0,0,0,0,1],size=size)
v = values>0


t=time.time()
numpy.apply_along_axis(find1,-1,v,1)
print('find1',time.time()-t)


t=time.time()
find2(v,1)
print('find2',time.time()-t)

输出

(1, 100000000)
('find1', 0.25300002098083496)
('find2', 0.2780001163482666)
(10000, 10000)
('find1', 0.46200013160705566)
('find2', 0.27300000190734863)
(1000000, 100)
('find1', 20.98099994659424)
('find2', 0.3040001392364502)
(10000000, 10)
('find1', 206.7590000629425)
('find2', 0.4830000400543213)

That said, a find written in C would be at least a little faster than either of these approaches

小开

虽然对你来说已经太迟了，但是为了以后的参考: 在 numba 实现之前，使用 numba (1)是最简单的方法。如果您使用的是 anaconda python 发行版，那么它应该已经被安装了。 The code will be compiled so it will be fast.

@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

and then:

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

这个怎么样

import numpy as np
np.amin(np.where(array==item))
小开

我为几种方法做了一个基准测试:

  • argwhere
  • 问题中的 nonzero
  • .tostring()是@Rob Reilink 的回答
  • 巨蟒循环
  • Fortran loop

Python和 Fortran代码是可用的。我跳过了那些没有希望的代码，比如转换成一个列表。

对数尺度上的结果。X 轴是指针的位置(查找指针是否在数组的下方需要更长的时间) ; 最后一个值是指针不在数组中。Y 轴是找到它的时间。

benchmark results

The array had 1 million elements and tests were run 100 times. Results still fluctuate a bit, but the qualitative trend is clear: Python and f2py quit at the first element so they scale differently. Python gets too slow if the needle is not in the first 1%, whereas f2py is fast (but you need to compile it).

总而言之，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)
小开

@ tal 已经提供了一个 numba函数来查找第一个索引，但这只适用于1D 数组。使用 np.ndenumerate，您还可以在任意多维数组中找到第一个索引:

from numba import njit
import numpy as np


@njit
def index(array, item):
for idx, val in np.ndenumerate(array):
if val == item:
return idx
return None

例子:

>>> arr = np.arange(9).reshape(3,3)
>>> index(arr, 3)
(1, 0)

计时显示，它的性能类似于 总数解决方案:

arr = np.arange(100000)
%timeit index(arr, 5)           # 1000000 loops, best of 3: 1.88 µs per loop
%timeit find_first(5, arr)      # 1000000 loops, best of 3: 1.7 µs per loop


%timeit index(arr, 99999)       # 10000 loops, best of 3: 118 µs per loop
%timeit find_first(99999, arr)  # 10000 loops, best of 3: 96 µs per loop
小开

作为一个长期的 matlab 用户，我一直在寻找一个有效的解决方案来解决这个问题已经有一段时间了。最后，由于讨论了 线中的一个命题，我试图提出一个解决方案，实现一个类似于建议的 给你的 API，目前只支持1D 数组。

You would use it like this

import numpy as np
import utils_find_1st as utf1st
array = np.arange(100000)
item = 1000
ind = utf1st.find_1st(array, item, utf1st.cmp_larger_eq)

支持的条件运算符是: cmp _ equals，cmp _ not _ equals，cmp _ large，cmp _ small，cmp _ large _ eq，cmp _ small _ eq。

你可在此找到资料来源、基准及其他详情:

Https://pypi.python.org/pypi?name=py_find_1st&:action=display

为了在我们的团队中使用(linux 和 macos 上的 anaconda) ，我已经做了一个可以简化安装的 anaconda 安装程序，您可以按照这里的描述使用它

Https://anaconda.org/roebel/py_find_1st

小开

这个问题可以通过以块的形式处理数组来有效地解决:

def find_first(x):
idx, step = 0, 32
while idx < x.size:
nz, = x[idx: idx + step].nonzero()
if len(nz): # found non-zero, return it
return nz[0] + idx
# move to the next chunk, increase step
idx += step
step = min(9600, step + step // 2)
return -1

数组以大小为 step的块处理。step的步长越长，处理零点阵列的速度越快(最坏的情况)。它越小，数组的处理速度越快，非零开始。诀窍是从一个小的 step开始，然后以指数级增长。此外，由于收益有限，没有必要将其增加到某个阈值以上。

我将该解决方案与纯 ndarary.nonzero 和 numba 解决方案与1000万个浮点数组进行了比较。

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


def find_first(x):
idx, step = 0, 32
while idx < x.size:
nz, = x[idx: idx + step].nonzero()
if len(nz):
return nz[0] + idx
idx += step
step = min(9600, step + step // 2)
return -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 54.733994480002366 ms
find_first 0.0013148509997336078 ms
find_first_numba 0.0002839310000126716 ms
---- LAST ----
ndarray.nonzero 54.56336712999928 ms
find_first 25.38929685000312 ms
find_first_numba 8.022820680002951 ms
---- NONE ----
ndarray.nonzero 24.13432420999925 ms
find_first 25.345200140000088 ms
find_first_numba 8.154927100003988 ms
---- ALL ----
ndarray.nonzero 55.753537260002304 ms
find_first 0.0014760300018679118 ms
find_first_numba 0.0004358099977253005 ms

纯 ndarray.nonzero绝对是比较宽松的。在最佳情况下，numba 解决方案的速度大约是5倍。在最坏的情况下，它的速度大约是3倍。

小开

如果你正在寻找第一个非零元素，你可以使用以下技巧:

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

这是一种 非常快“ numpy-Pure”解决方案，但在下面讨论的某些情况下失败了。

该解决方案利用了这样一个事实，即数值类型的几乎所有零表示都由 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

溶液中 faster含量比 numba 高33% ，为“ numpy 纯”溶液。

缺点:

  • 不适用于像 object这样麻木的可接受类型
  • 失败的负零，偶尔出现在 float或 double计算

提交答案