如何使用pandas计算分组内行差值并新增为DataFrame附加列

需求描述 #

需要对如下测试DataFrame按id列分组：

id  a  b  c
0   1  1  6  2
1   1  2  5  2
2   2  3  4  2
3   2  4  3  2
4   3  5  2  2
5   3  6  1  2

计算同一分组内a/b/c列的行之间的差值，将差值作为附加列添加到原DataFrame中，最终输出结构如下：

id  a  b  c  a_diff  b_diff  c_diff
0   1  1  6  2    -1.0     1.0     0.0
1   1  2  5  2     1.0    -1.0     0.0
2   2  3  4  2    -1.0     1.0     0.0
3   2  4  3  2     1.0    -1.0     0.0
4   3  5  2  2    -1.0     1.0     0.0
5   3  6  1  2     1.0    -1.0     0.0

本次使用的测试数据构造代码：

import pandas as pd
df = pd.DataFrame({'id': [1,1,2,2,3,3], 'a': [1,2,3,4,5,6],'b': [6,5,4,3,2,1], 'c': [2,2,2,2,2,2]})

实现代码 #

观察目标结果的差值规律：每个分组共2行，第一行差值为当前行减同组下一行，第二行差值为当前行减同组上一行，实现逻辑如下：

# 按id分组计算后行减前行的差值，列名追加_diff后缀
diff_cols = df.groupby('id')[['a', 'b', 'c']].diff().add_suffix('_diff')
# 填充每个分组首行的空值：取同组第二行差值的相反数
diff_cols = diff_cols.groupby(df['id']).apply(lambda group: group.fillna(-group.dropna().iloc[0])).reset_index(drop=True)
# 拼接原表和差值列得到最终结果
df_result = pd.concat([df, diff_cols], axis=1)

执行后打印df_result即可得到和要求完全一致的输出。如果分组内存在更多行，可根据实际差值计算规则调整填充逻辑即可。

内容的提问来源于stack exchange，提问作者JFG123