如何在Python中按id分组，按薪资差值升序排序DataFrame行？

问题描述 #

现有如下Pandas DataFrame，每个id对应的Salary重复出现，需要将每个id组内的行按照Salary与Reference_Salary的差值从小到大排序，得到目标结构的DataFrame。

给定DataFrame创建代码 #

# 导入所需库
import pandas as pd
import matplotlib.pyplot as plt

# 创建数据集
data_set = {'id': [1,1,1,1,1,2,2,2,2,2,3,3,3,3,3],
            'Salary': [16,16,16,16,16,42,42,42,42,42,28,28,28,28,28],
            'Reference_id': [1,2,3,4,5,1,2,3,4,5,1,2,3,4,5],
            'Reference_Salary': [10,20,30,40,50,10,20,30,40,50,10,20,30,40,50],
            'Expenditure': [100,210,320,430,540,90,200,310,420,530,80,190,300,410,520],
            'Leaves': [17.5,45.,72.5,100.,127.5,7.5,35.,62.5,90.,117.5,2.5,25.,52.5,80.,107.5]}

# 转换为DataFrame
df = pd.DataFrame(data_set)
print("\n df = \n",df)

解决方案 #

以下两种方法均可实现需求，其中方法一性能更优，适合大数据集场景：

方法一：计算差值列后排序（推荐） #

先生成差值列，再按id和差值列升序排序，最后可选择删除临时差值列：

# 计算Salary与Reference_Salary的差值
df['salary_diff'] = df['Salary'] - df['Reference_Salary']

# 按id分组，再按差值升序排序
sorted_df = df.sort_values(by=['id', 'salary_diff'], ascending=True)

# 可选：删除临时差值列
sorted_df = sorted_df.drop(columns=['salary_diff'])

print("\n 排序后的DataFrame = \n", sorted_df)

方法二：分组后应用排序 #

通过groupby分组后对每组数据直接按差值排序，无需额外生成列：

# 按id分组，每组内按Salary与Reference_Salary的差值升序排序
sorted_df = df.groupby('id', group_keys=False).apply(
    lambda x: x.sort_values(by=x['Salary'] - x['Reference_Salary'], ascending=True)
)

print("\n 排序后的DataFrame = \n", sorted_df)

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