时间戳

首页 > 解决方案 > 在数组中找到下一个最大值的更快方法

python - 在数组中找到下一个最大值的更快方法

问题描述

我正在尝试改善我的数据帧代码的运行时间。我的思路错了吗？

我有以下代码来查找第 1 列中的第一个值，该值高于value并具有比它更高的索引（index_value=n）

new_index=(df[n:,1] > value).argmax()

我的问题是：argmax() 参数将构建一个包含 True 和 False 的完整列表，然后，它才会找到第一次出现并返回我预期的索引。

有没有办法改进这段代码？即在找到第一个列表后停止构建列表True。

标签： pythonarrayspandasnumpy

解决方案

本来不打算发帖的。我期待 numba 在所有情况下都能获胜，但这并不是注定的。对提议的解决方案进行了一些基准测试，结果有些有趣，因此在此处发布。我将使用数组数据来保持简单。

# Proposed solutions
import numpy as np
from numba import njit

# @piRSquared's soln
@njit
def find_first_gt(a, n, value):
    while a[n] <= value:
        n += 1
    return n

# @Ehsan's soln
def numpy_argmax(a, n , value):
    return np.argmax(a[n:] > value)

使用benchit包（几个基准测试工具打包在一起；免责声明：我是它的作者）对建议的解决方案进行基准测试。

计时和加速 -

# Benchmark
a = np.arange(1000_000)
n = 0

import benchit
funcs = [find_first_gt, numpy_argmax]
vs = np.linspace(0, len(a)-1, num=20, endpoint=True).astype(int)
inputs = [(a,0,v) for v in vs]
t = benchit.timings(funcs, inputs, multivar=True, input_name='Position of value')
t.plot(logy=False, logx=False, savepath='plot.png')
t.speedups(ref_func_by_index=1).plot('Speedup_with_numba.png')

在此处输入图像描述

在此处输入图像描述

如果您对确切的加速数字感兴趣 -

In [12]: t.speedups(ref_func_by_index=1)
Out[12]: 
Functions          find_first_gt  Ref:numpy_argmax
Position of value                                 
0                    2103.548010               1.0
52631                  22.053699               1.0
105263                 11.109615               1.0
157894                  7.541725               1.0
210526                  5.640514               1.0
263157                  4.407300               1.0
315789                  3.642989               1.0
368420                  3.028726               1.0
421052                  2.543713               1.0
473683                  2.201336               1.0
526315                  1.931540               1.0
578946                  1.692138               1.0
631578                  1.536912               1.0
684209                  1.455065               1.0
736841                  1.357728               1.0
789472                  1.248716               1.0
842104                  1.176199               1.0
894735                  1.062174               1.0
947367                  1.043791               1.0
999999                  0.983419               1.0

结论：在几乎所有情况下，numba都做得很好，除非您知道value它位于最远端，或者 numba 缓存方案对您不利。

推荐阅读