Сгруппируйте, примените функцию к каждой строке со смещением и создайте новый столбец

Question

Я хочу сгруппировать по id, применить функцию к данным и создать новый столбец с результатами.Кажется, должен быть более быстрый / более эффективный способ сделать это, чем передавать данные в функцию, вносить изменения и возвращать данные.Вот пример.

Пример

dat = pd.DataFrame({'id': ['a', 'a', 'a', 'b', 'b', 'b'], 'x': [4, 8, 12, 25, 30, 50]})

def my_func(data):
    data['diff'] = (data['x'] - data['x'].shift(1, fill_value=data['x'].iat[0]))
    return data

dat.groupby('id').apply(my_func)

Вывод

> print(dat)


id  x   diff
0   a   4   0
1   a   8   4
2   a   12  4
3   b   25  0
4   b   30  5
5   b   50  20

Есть ли более эффективный способсделать это?

Answer 1

Вы можете использовать .groupby.diff() для этого и после этого заполнить NaN нулями, как показано ниже:

dat['diff'] = dat.groupby('id').x.diff().fillna(0)

print(dat)
  id   x  diff
0  a   4   0.0
1  a   8   4.0
2  a  12   4.0
3  b  25   0.0
4  b  30   5.0
5  b  50  20.0