Pandas高效实现双列分组求和后取每组最大值对应SUB_ID

我来帮你解决这个适合百万级数据的高效处理问题，咱们一步步来优化：

步骤1：优化分组绝对值求和 #

你原来的apply写法虽然能得到结果，但对于百万级数据来说，用pandas内置的向量化聚合会大幅降低开销——避免逐组的Python lambda调用，改用底层C实现的操作：

import pandas as pd
import numpy as np

# 构造测试数据
data = {
    'ID': [101,101,101,101,101,102,102,102,102],
    'SUB_ID': [1,1,1,2,2,3,3,4,4],
    'AMOUNT': [50,-10,-20,30,20,10,-10,10,10]
}
df = pd.DataFrame(data)

# 高效计算每个(ID, SUB_ID)组的绝对值和
sum_df = df.assign(ABS_AMOUNT=df['AMOUNT'].abs()) \
           .groupby(['ID', 'SUB_ID'])['ABS_AMOUNT'] \
           .sum() \
           .reset_index(name='ABS_SUM')

运行后sum_df的结果如下：

步骤2：提取每个ID组内最大值对应的SUB_ID #

这里推荐两种完全无循环的向量化方案，都适合百万级数据：

方案1：用idxmax定位最大值行 #

idxmax会直接找到每个ID组内ABS_SUM最大的行索引，效率极高：

# 获取每个ID组内最大值对应的SUB_ID，转成列表
result = sum_df.loc[sum_df.groupby('ID')['ABS_SUM'].idxmax(), 'SUB_ID'].tolist()

测试数据下结果为[1, 3]（ID102中取第一个出现的最大值对应的SUB_ID）。

方案2：排序后去重（灵活选择并列值） #

如果需要灵活选择并列最大值的任意一个（比如取最后一个），可以先排序再去重：

# 按ID升序、ABS_SUM降序排序，保留每个ID的第一行
result = sum_df.sort_values(['ID', 'ABS_SUM'], ascending=[True, False]) \
               .drop_duplicates('ID') \
               ['SUB_ID'] \
               .tolist()

# 若要取并列最大值的最后一个，修改排序和去重参数
result = sum_df.sort_values(['ID', 'ABS_SUM'], ascending=[True, True]) \
               .drop_duplicates('ID', keep='last') \
               ['SUB_ID'] \
               .tolist()

为什么这些方法高效？ #

所有操作都是基于pandas底层C语言实现的向量化运算，完全避免了Python级别的循环，处理百万行数据时性能会比自定义apply或手写循环快几个数量级，不会出现内存或速度瓶颈。

如果要直接基于你原来生成的MultiIndex Series处理，只需先转成DataFrame再用上述方法：

# 基于你原来的df1处理
df1 = df.groupby(['ID','SUB_ID'])['AMOUNT'].apply(lambda x: np.sum(np.absolute(x)))
sum_df = df1.reset_index(name='ABS_SUM')
result = sum_df.loc[sum_df.groupby('ID')['ABS_SUM'].idxmax(), 'SUB_ID'].tolist()

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