Pandas - 将 groupby() 与 rolling() 和 apply()

Question

在 Pandas 中为大型数据帧（1500682 行）调用 groupby 以及滚动和应用函数时，我的性能非常缓慢。我正在尝试获得具有不同权重的滚动移动平均线。

运行缓慢的代码部分是：

df['rolling'] = df.groupby('i2')['x'].rolling(3).apply(lambda x: x[-3]*0.1+x[-2]*0.9).reset_index(level=0, drop=True).reindex(df.index)

完整的代码（带有数据）是：

import pandas as pd

from random import randint

# data (it takes some time to create [less than 1 minute in my computer])

data1   = [[[[randint(0, 100) for i in range(randint(1, 2))] for i in range(randint(1, 3))] for i in range(5000)] for i in range(100)]

data2   = pd.DataFrame(

    [

        (i1, i2, i3, i4, x4)

        for (i1, x1) in enumerate(data1)

        for (i2, x2) in enumerate(x1)

        for (i3, x3) in enumerate(x2)

        for (i4, x4) in enumerate(x3)

    ],

    columns = ['i1', 'i2', 'i3', 'i4', 'x']

)

data2.drop(['i3', 'i4'], axis=1, inplace = True)

df   = data2.set_index(['i1', 'i2']).sort_index()

## conflicting part of the code ##

df['rolling'] = df.groupby('i2')['x'].rolling(3).apply(lambda x: x[-3]*0.1+x[-2]*0.9).reset_index(level=0, drop=True).reindex(df.index)

如果您能详细说明代码以使其更高效并更快地执行，我将不胜感激。

Answer 1

如果我理解正确，您可以尝试：grp=df.groupby('i2')['x'] df['rolling']=grp.shift(2).mul(0.1).add(grp.shift(1).mul(0.9))现在详细说明：为什么不.apply(...)：我什么时候应该在我的代码中使用 pandas apply()？你应该做的是使用任何利用矢量化操作的东西。我在这里做了一些更详细的解释：https://stackoverflow.com/a/60029108/11610186