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

numba

有了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级别的优化来提高性能。一个 常规的for无环路@njit装饰通常会 慢 比你已经尝试了在条件满足后期的情况下的方法。

对于Pandas数值系列df['data']，您可以简单地将NumPy表示提供给JIT编译的函数：

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


概括

由于numba允许以函数作为参数，并且假设传递的函数也可以JIT编译，则可以找到一种方法来计算第 n 个索引，其中满足任意条件的条件func。

@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)


对于第三个 最后的 值，可以喂相反，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