如何加快熊猫滚动argmax？答案

【问题标题】：how to speed up pandas rolling argmax?如何加快熊猫滚动argmax？
【发布时间】：2021-10-17 08:21:31
【问题描述】：

我有一个很长的熊猫系列。

例如：

In [1]: import pandas as pd

In [2]: a = pd.DataFrame(range(100000000))

我想在滚动 10 列表中获取最大元素的索引。

我试过了：

In [8]: a[0].rolling(10).apply(lambda x: x.idxmax())

和

In [9]: a[0].rolling(10).apply(np.argmax)

两者都运行太慢，有什么方法可以让这个工作在一个稳定的时间内完成吗？

【问题讨论】：

【解决方案1】：

使用numba引擎进行快速计算，与不使用numba相比，您将获得更快的结果，除了numba还有cython优化。

Numba 与 numpy 配合得很好，pandas 与 numpy 配合得很好，有时你可以做 a.to_numpy 并在 numpy 中使用或不使用 numba 进行计算并重新分配给 pandas

import pandas as pd

a = pd.DataFrame(range(10000000))
a.rolling(10).apply(np.argmax, engine='numba', raw=True)

【讨论】：

嗨，eroot，你知道如何加速 argsort 吗？ a.rolling(10).apply(lambda x: np.argsort(x)[-1], engine='numba', raw=True) 似乎很慢。
我对此不太确定，但是当尝试使用未优化的 lambda 自定义函数时，它往往会变慢。我刚刚开始进行性能优化，所以不太确定，但为了速度，我总是更喜欢内置的东西
当我知道更多确切原因时，我会更新。事实上，你的问题可能是一个关于为什么 lamda 速度慢的新问题。甚至 lambda x: np.argmax(x) 实际上比直接使用 np.argmax 慢。可能它与 lambda 是一个调用另一个函数的函数有关，所以 lambda 设置了一个更多的堆栈？？
argsort 需要排序，这就是为什么它很慢并且需要时间