如何用Python Pandas高效按批次计算加权平均值（规避除零）

问题背景 #

现有10万行规模的Pandas DataFrame，按batch字段划分批次，数据结构如下：

batch    b       c
 A      10       2     
 A      20       4     
 A      30       6     
 B      5        1     
 B      10       2   
 B      15       3   
 B      20       4    
 .       .       .     
 .       .       . 
 .       .       .     

需按batch分组，每组计算加权平均值（分子为b*c的求和，分母为c的求和），生成目标格式DataFrame：

batch   new_value
 A        23.3 
 B        15.0
 .         .
 .         .
 .         .

已知c列存在零值，需避免除零异常；当前实现为新增b*c列后分组求和再相除，寻求更高效的内存与性能优化方案。

优化实现方案 #

方案1：无中间列的向量化聚合（最优性能） #

直接在分组聚合阶段完成分子与分母的计算，无需创建额外列，大幅降低内存开销，同时利用Pandas向量化操作保证速度：

import pandas as pd
import numpy as np

# 假设原始数据集为df
result = df.groupby('batch').agg(
    numerator=('b', lambda x: (x * df.loc[x.index, 'c']).sum()),
    denominator=('c', 'sum')
).assign(
    new_value=lambda df_agg: np.divide(
        df_agg['numerator'], 
        df_agg['denominator'], 
        out=np.full_like(df_agg['numerator'], np.nan),  # 分母为0时填充NaN
        where=df_agg['denominator'] != 0
    )
).reset_index()[['batch', 'new_value']]

关键细节：

• 聚合阶段直接计算b*c的总和，避免存储中间列（10万行数据可节省约一列的内存）
• 用np.divide的where参数精准规避除零场景，分母为0时结果设为NaN（可按需改为0或其他默认值）
• 最后仅保留目标列，生成符合要求的结构

方案2：简洁的apply实现（代码短，性能略逊） #

若优先追求代码简洁，可使用groupby.apply，但大数据量下性能略低于方案1：

result = df.groupby('batch').apply(
    lambda group: np.divide(
        (group['b'] * group['c']).sum(), 
        group['c'].sum(), 
        out=np.array([np.nan]), 
        where=group['c'].sum() != 0
    )[0]
).reset_index(name='new_value')

性能对比 #

• 原始方案（新增中间列）：额外存储b*c列，内存开销增加约一列的容量，聚合阶段需读取两列数据
• 优化方案：无中间列存储，聚合时直接计算，内存效率提升明显；方案1的向量化操作速度优于方案2与原始方案

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