如何用向量化替代groupby后的for循环加速Pandas代码？

本人刚接触pandas、dataframe等工具。目前需处理展开后达数亿行的数据集，希望优化代码以提升运行效率。我通过.groupby('ORG_INDEX')得到数十万个分组，需对每个分组执行简单运算并新增列，但当前的for循环函数耗时过长。听闻向量化可替代for循环，但找不到适配的示例，现附上原函数：

def add_data_df(df1,df2):
  df1_group = df1.copy().groupby('ORG_INDEX', sort=False)
  sums=[]
  min_gfa=[]
  max_gfa=[]
  min_uni=[]
  max_uni=[]
  for group in df1_group:
    groupdf = group[1]
    sum_a = np.sum(groupdf['F_AREA'])
    sums.append(sum_a)

    groupdf['MINGFA'] = (groupdf['MIN_FLR'] * groupdf['F_AREA']) + (groupdf['MIN_FLR2'] * groupdf['F_AREA2'])
    groupdf['MAXGFA'] = (groupdf['MAX_FLR'] * groupdf['F_AREA']) + (groupdf['MAX_FLR2'] * groupdf['F_AREA2'])
    groupdf['MINUNITS'] = (groupdf['MIN_FLR'] * groupdf['UNITS_F']) + (groupdf['MIN_FLR2'] * groupdf['UNITS_F2'])
    groupdf['MAXUNITS'] = (groupdf['MAX_FLR'] * groupdf['UNITS_F']) + (groupdf['MAX_FLR2'] * groupdf['UNITS_F2'])

    min_gfa.append(np.sum(groupdf['MINGFA']))
    max_gfa.append(np.sum(groupdf['MAXGFA']))
    min_uni.append(np.sum(groupdf['MINUNITS']))
    max_uni.append(np.sum(groupdf['MAXUNITS']))

  df2['SUM_AREA']=sums
  df2['MIN_GFA']=min_gfa
  df2['MAX_GFA']=max_gfa
  df2['MIN_UNITS']=min_uni
  df2['MAX_UNITS']=max_uni

寻求将其改为向量化实现的方法。

原代码的核心问题是手动遍历分组，完全没利用pandas的向量化运算能力。以下是优化后的实现，全程用向量化操作替代循环，效率会提升几个数量级：

步骤1：计算中间列（逐行向量化运算） #

四个中间列MINGFA、MAXGFA、MINUNITS、MAXUNITS的计算是逐行独立的，不需要分组，直接在原df1上计算即可：

# 直接在df1上计算中间列，无需分组
df1['MINGFA'] = df1['MIN_FLR'] * df1['F_AREA'] + df1['MIN_FLR2'] * df1['F_AREA2']
df1['MAXGFA'] = df1['MAX_FLR'] * df1['F_AREA'] + df1['MAX_FLR2'] * df1['F_AREA2']
df1['MINUNITS'] = df1['MIN_FLR'] * df1['UNITS_F'] + df1['MIN_FLR2'] * df1['UNITS_F2']
df1['MAXUNITS'] = df1['MAX_FLR'] * df1['UNITS_F'] + df1['MAX_FLR2'] * df1['UNITS_F2']

步骤2：分组聚合统计量 #

利用groupby.agg()一次性计算所有需要的聚合值，pandas会自动用向量化方式处理分组：

# 按ORG_INDEX分组，聚合所有需要的统计量
grouped_stats = df1.groupby('ORG_INDEX', sort=False).agg(
    SUM_AREA=('F_AREA', 'sum'),
    MIN_GFA=('MINGFA', 'sum'),
    MAX_GFA=('MAXGFA', 'sum'),
    MIN_UNITS=('MINUNITS', 'sum'),
    MAX_UNITS=('MAXUNITS', 'sum')
).reset_index(drop=True)  # 重置索引，确保和df2的行顺序匹配

步骤3：将聚合结果赋值给df2 #

直接把聚合后的DataFrame的列批量赋值给df2即可：

def add_data_df(df1, df2):
    # 若不想修改原df1，可先复制一份
    df1 = df1.copy()
    
    # 步骤1：计算中间列
    df1['MINGFA'] = df1['MIN_FLR'] * df1['F_AREA'] + df1['MIN_FLR2'] * df1['F_AREA2']
    df1['MAXGFA'] = df1['MAX_FLR'] * df1['F_AREA'] + df1['MAX_FLR2'] * df1['F_AREA2']
    df1['MINUNITS'] = df1['MIN_FLR'] * df1['UNITS_F'] + df1['MIN_FLR2'] * df1['UNITS_F2']
    df1['MAXUNITS'] = df1['MAX_FLR'] * df1['UNITS_F'] + df1['MAX_FLR2'] * df1['UNITS_F2']
    
    # 步骤2：分组聚合
    grouped_stats = df1.groupby('ORG_INDEX', sort=False).agg(
        SUM_AREA=('F_AREA', 'sum'),
        MIN_GFA=('MINGFA', 'sum'),
        MAX_GFA=('MAXGFA', 'sum'),
        MIN_UNITS=('MINUNITS', 'sum'),
        MAX_UNITS=('MAXUNITS', 'sum')
    ).reset_index(drop=True)
    
    # 步骤3：批量赋值给df2
    df2[['SUM_AREA', 'MIN_GFA', 'MAX_GFA', 'MIN_UNITS', 'MAX_UNITS']] = grouped_stats
    
    return df2

额外优化建议 #

• 内存优化：若df1内存紧张，计算完中间列后可直接聚合，再删除中间列释放内存：df1.drop(['MINGFA', 'MAXGFA', 'MINUNITS', 'MAXUNITS'], axis=1, inplace=True)。
• 顺序一致性：因为设置了sort=False，分组结果的顺序和df1中ORG_INDEX首次出现的顺序一致，确保df2的行顺序与之对应即可避免赋值错误。

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