Рекомбинировать групповую сумму с исходным pandas DataFrame

У меня есть 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, который содержит « столбец с'. Я дошел до:

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 для каждой группы в массив, заархивировать его с помощью аналогичного массива индексов и т. д., но это кажется очень хакерским. Есть ли встроенный метод панд, который мне здесь не хватает?


python pandas group-by rolling-sum
person user1684046    schedule 28.11.2016    source источник

Ответы (1)


arrow_upward
1
arrow_downward

Если использование groupby является обязательным, вы можете использовать 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

введите здесь описание изображения

person Nickil Maveli    schedule 28.11.2016
comment
вы можете сделать это напрямую: FYI pandas.pydata.org/pandas-docs/stable/ (хотя я понимаю, что это не задокументировано, кроме как в Whatsnew) - person Jeff; 29.11.2016
comment
и если кто-то хочет улучшить документы: github.com/pandas-dev/pandas/issues /14759 - person Jeff; 29.11.2016