如何加快 pandas 滚动 argmax？

Question

我有一个很长的pandas系列。

例如：

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)

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

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

PS：查看我在 上的回答，以实现从 pandas 到 numpy

的高效转换