如何优化Pandas groupby+idxmax取分组最大值的代码性能

需求说明 #

• 待处理数据为Airbnb房源价格数据集，包含neighbourhood group（社区组）、neighbourhood（社区）、room_type（房源类型）、收益等字段
• 目标输出：每个社区仅保留1条平均收益最高的房源类型及对应价格数据
• 原有写法的缺陷：依赖显式Python循环做索引筛选，运行效率低，代码冗余度高，不符合简洁的编码规范

原有实现代码 #

df = airbnb.loc[airbnb['neighbourhood'].isin(list(combined['neighbourhood']))].groupby(['neighbourhood','room_type']).mean().sort_values(by=['Revenues'],ascending=False)['Revenues'].reset_index().sort_values(by=['neighbourhood','room_type']).reset_index()
tmp_list = []
test = pd.DataFrame(df.groupby(['neighbourhood','room_type'])['Revenues'].idxmax(axis=0)).reset_index()
# 回写收益值用于后续筛选
test['Actual_Revenues'] = df['Revenues']
# 循环遍历每个社区取最大值索引
for i in neighbourhood_list:
    print(test[test['neighbourhood']==i])
    max_index_list = test[test['neighbourhood']==i].sort_values(by='Actual_Revenues', ascending=False).head(1)['Revenues']
    print(max_index_list)
    tmp_list.append(list(max_index_list))
    
tmp_list = list(np.concatenate(tmp_list).flat)
df[df.index.isin(tmp_list)]

优化后实现 #

直接使用Pandas内置的向量化方法完成分组筛选，完全取消显式循环，代码更简洁、运行效率更高：

# 计算各社区+各房源类型维度的平均收益，可按需保留需要的字段
avg_revenue_df = airbnb[
    airbnb['neighbourhood'].isin(combined['neighbourhood'])
].groupby(
    ['neighbourhood group', 'neighbourhood', 'room_type'], as_index=False
)['Revenues'].mean()

# 按收益降序排序后，每个社区保留第一条（即最高收益）记录
final_result = avg_revenue_df.sort_values(
    by='Revenues', ascending=False, ignore_index=True
).drop_duplicates(subset='neighbourhood', keep='first')

特殊场景适配：如果同一社区存在多个房源类型平均收益完全相同、需要保留所有并列第一的记录，可以将第二步替换为以下代码：

final_result = avg_revenue_df.groupby('neighbourhood', as_index=False).apply(
    lambda group: group[group['Revenues'] == group['Revenues'].max()]
).reset_index(drop=True)

优化收益：所有运算均走Pandas底层向量化逻辑，没有Python层循环开销，数据量越大性能优势越明显，代码可读性也显著提升。

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