python返回数组索引_python-有效地返回数组中第一个值满足条件的索引

我需要在满足条件的1d NumPy数组或Pandas数值系列中找到第一个值的索引.数组很大,索引可能在数组的开始或结尾附近,或者可能根本不满足条件.我无法提前告诉您哪种可能性更大.如果不满足条件,则返回值应为-1.我考虑了几种方法.

尝试1

# func(arr) returns a Boolean array

idx = next(iter(np.where(func(arr))[0]), -1)

但这通常太慢,因为func(arr)在整个数组上应用矢量化函数,而不是在满足条件时停止.具体来说,在数组开始附近满足条件时,这很昂贵.

尝试2

np.argmax的速度稍快一些,但无法确定何时从不满足条件：

np.random.seed(0)

arr = np.random.rand(10**7)

assert next(iter(np.where(arr > 0.999999)[0]), -1) == np.argmax(arr > 0.999999)

%timeit next(iter(np.where(arr > 0.999999)[0]), -1) # 21.2 ms

%timeit np.argmax(arr > 0.999999) # 17.7 ms

np.argmax(arr> 1.0)返回0,即不满足条件的实例.

尝试3

# func(arr) returns a Boolean scalar

idx = next((idx for idx, val in enumerate(arr) if func(arr)), -1)

但这在数组末尾附近满足条件时太慢了.大概是因为生成器表达式具有来自大量__next__调用的昂贵开销.

这是否总是一种折衷方案？对于通用函数,有没有办法有效地提取第一个索引？

标杆管理

对于基准测试,假设func在值大于给定常量时找到索引：

# Python 3.6.5, NumPy 1.14.3, Numba 0.38.0

import numpy as np

np.random.seed(0)

arr = np.random.rand(10**7)

m = 0.9

n = 0.999999

# Start of array benchmark

%timeit next(iter(np.where(arr > m)[0]), -1) # 43.5 ms

%timeit next((idx for idx, val in enumerate(arr) if val > m), -1) # 2.5 µs

# End of array benchmark

%timeit next(iter(np.where(arr > n)[0]), -1) # 21.4 ms

%timeit next((idx for idx, val in enumerate(arr) if val > n), -1) # 39.2 ms

解决方法:

麻巴

使用numba可以优化两种情况.从语法上讲,您只需要构造一个带有简单for循环的函数：

from numba import njit

@njit

def get_first_index_nb(A, k):

for i in range(len(A)):

if A[i] > k:

return i

return -1

idx = get_first_index_nb(A, 0.9)

Numba通过JIT(“及时”)编译代码并利用CPU-level optimisations来提高性能.不带@njit装饰器的常规for循环通常比您在条件较晚满足的情况下尝试过的方法要慢.

对于熊猫数字系列df [‘data’],您可以将NumPy表示形式简单地馈送到JIT编译的函数中：

idx = get_first_index_nb(df['data'].values, 0.9)

概括

由于numba允许functions as arguments,并且假设传递的函数也可以JIT编译,则可以找到一种方法,该方法可以计算第n个索引,其中满足任意函数的条件.

@njit

def get_nth_index_count(A, func, count):

c = 0

for i in range(len(A)):

if func(A[i]):

c += 1

if c == count:

return i

return -1

@njit

def func(val):

return val > 0.9

# get index of 3rd value where func evaluates to True

idx = get_nth_index_count(arr, func, 3)

对于倒数第3个值,您可以输入相反的值arr [::-1],并取反len(arr)-1的结果,而-1是解决0索引所必需的.

绩效基准

# Python 3.6.5, NumPy 1.14.3, Numba 0.38.0

np.random.seed(0)

arr = np.random.rand(10**7)

m = 0.9

n = 0.999999

@njit

def get_first_index_nb(A, k):

for i in range(len(A)):

if A[i] > k:

return i

return -1

def get_first_index_np(A, k):

for i in range(len(A)):

if A[i] > k:

return i

return -1

%timeit get_first_index_nb(arr, m) # 375 ns

%timeit get_first_index_np(arr, m) # 2.71 µs

%timeit next(iter(np.where(arr > m)[0]), -1) # 43.5 ms

%timeit next((idx for idx, val in enumerate(arr) if val > m), -1) # 2.5 µs

%timeit get_first_index_nb(arr, n) # 204 µs

%timeit get_first_index_np(arr, n) # 44.8 ms

%timeit next(iter(np.where(arr > n)[0]), -1) # 21.4 ms

%timeit next((idx for idx, val in enumerate(arr) if val > n), -1) # 39.2 ms

标签：performance,pandas,arrays,python,numpy

来源： https://codeday.me/bug/20191024/1924444.html