панды dataframe окно прокатки с GroupBy

Question

я могу добавить новый столбец c, который является суммой двух последних значений в b, как показано ниже ...

df['c'] = df.b.rolling(window = 2).sum().shift() 

df 
    a b  c 
0 1 3 NaN 
1 1 0 NaN 
2 0 6 3.0 
3 1 0 6.0 
4 0 0 6.0 
5 1 7 0.0 
6 0 0 7.0 
7 0 7 7.0 
8 1 4 7.0 
9 1 2 11.0 

... однако, что если я хочу группа по a сначала? Например. Я могу это сделать:

df['c'] = df.groupby(['a'])['b'].shift(1) + df.groupby(['a'])['b'].shift(2)

Есть ли более элегантный способ для суммирования большого числа сдвигов (1, 2, ... п) на группе?

Answer 1

f = lambda x: x.rolling(2).sum().shift() 
df['c'] = df.groupby('a').b.apply(f) 

df