通过与原始pandas DataFrame的rollling sum重新组合

Question

我有一个形式的pandas DataFrame：

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附加到一个数组，用一个类似的索引数组拉链等等，但这看起来非常h​​acky。我在这里找不到内置的熊猫方法吗？

Answer 1

如果必须使用groupby，则可以使用groupby.apply一次性计算所有内容：

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

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

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

两者都给你： -

df