python - 将 groupby 滚动总和与原始 pandas DataFrame 重新组合

Question

我有一个形式的熊猫数据框：

import pandas as pd

df = pd.DataFrame({
    'a': [1,2,3,4,5,6],
    'b': [0,1,0,1,0,1]
})

我想按'b'的值对数据进行分组，并添加新列'c'，其中包含每个组的滚动总和'a'，然后我想将所有组重新组合成一个未分组的DataFrame，其中包含' c' 列。我已经做到了：

for i, group in df.groupby('b'):
    group['c'] = group.a.rolling(
        window=2,
        min_periods=1,
        center=False
    ).sum()

但是这种方法存在几个问题：

使用 for 循环对每个组进行操作感觉对于大型 DataFrame 来说会很慢（就像我的实际数据一样）
我找不到一种优雅的方式来保存每个组的“c”列并将其添加回原始 DataFrame。我可以将每个组的 c 附加到一个数组中，用一个类似的索引数组压缩它，等等，但这似乎很 hacky。我在这里缺少内置的 pandas 方法吗？

score 1 · Accepted Answer

如果 usinggroupby是必须的，那么您可以使用groupby.apply一次性计算所有内容：

df['c'] = df.groupby('b')['a'].apply(lambda x: x.rolling(2, min_periods=1).sum())

从开始v0.19.1，您可以直接调用rolling()/expanding()groupby 对象上的方法，如下所示：

df['c'] = df.groupby('b').rolling(2, min_periods=1)['a'].sum().sortlevel(1).values

两者都给你：-

df

1 回答 1