python - Pandas groupby diff 删除列

Question

我有一个这样的数据框:

d = {'id': ['101_i','101_e','102_i','102_e'], 1: [3, 4, 5, 7], 2: [5,9,10,11], 3: [8,4,3,7]}
df = pd.DataFrame(data=d)

我想减去所有具有相同前缀 ID 的行，即用 101_e 减去 101_i 行的所有值，反之亦然。我为此使用的代码是:

df['new_identifier'] = [x.upper().replace('E', '').replace('I','').replace('_','') for x in df['id']]
df = df.groupby('new_identifier')[df.columns[1:-1]].diff().dropna()

我得到这样的输出:

我看到我丢失了我创建的新列 new_identifier。有什么办法可以保留它吗？

Answer 1

您可以为已知类型(在本例中为 int 或 float)的列定义特定的聚合函数(在本例中为第 1、2 和 3 列的 np.diff())。

import numpy as np
df.groupby('new_identifier').agg({i: np.diff for i in range(1, 4)}).dropna()

结果:

                1  2  3
new_identifier         
101             1  4 -4
102             2  1  4