如何在Pandas中按Class分组实现加权平均成本(WAC)计算？

需求说明 #

现有包含Class、Rec、Iss、Price、Amt、Qtty字段的DataFrame，需生成带有WAC列的目标DataFrame。原代码仅适用于单一产品类，现需按Class分组后，对每个分组独立应用原有WAC计算逻辑。

输入输出示例 #

输入DataFrame #

Class   Rec  Iss   Price Amt  Qtty
 FD      10    0    5     50   10
 FD       0    2    4     -8    8
 RM       5    0    2     10    5
 RM       0    3    3     -9    2
 FD       0    2    2     -4    6
 PM       5    0    3     15    5
 PM       1    0    3      3    6
 RM       1    0    3      3    3
 FD       4    0    2      8    10

目标DataFrame #

Class   Rec  Iss   Price Amt  Qtty  WAC
 FD      10    0    5     50   10     5
 FD       0    2    4     -8    8     5
 RM       5    0    2     10    5     2
 RM       0    3    3     -9    2     2
 FD       0    2    2     -4    6     5
 PM       5    0    3     15    5     3
 PM       1    0    3      3    6     3
 RM       1    0    3      3    3   2.3
 FD       4    0    2      8   10   3.8

原单类代码 #

# 初始化首行平均价格
df.loc[0,'WAC'] = df.loc[0,'Price']

# 遍历剩余行计算WAC
for i in range(1, len(df)):
    # 无收货时继承前一行WAC
    if df.loc[i, 'Rec'] == 0:
        df.loc[i, 'WAC'] = df.loc[i-1, 'WAC']
    # 有收货时重新计算加权平均
    elif df.loc[i, 'Rec'] > 0:
        df.loc[i, 'WAC'] = (df.loc[i,'Amt'] + ((df.loc[i-1, 'WAC']) * (df.loc[i-1, 'Qtty']))) / (df.loc[i, 'Qtty'])
    # 收货为负时继承前一行WAC
    elif df.loc[i, 'Rec'] < 0:
        df.loc[i, 'WAC'] = df.loc[i-1, 'WAC']
    else:
        df.loc[i, 'WAC'] = 0
print(df)

分组实现方案 #

步骤1：定义分组计算WAC的函数 #

将原逻辑封装为函数，针对单个分组独立计算：

def calculate_wac(group):
    # 复制分组数据，避免修改原DataFrame
    group = group.copy()
    # 初始化分组首行的WAC为对应Price值
    first_idx = group.index[0]
    group.loc[first_idx, 'WAC'] = group.loc[first_idx, 'Price']
    
    # 遍历分组内剩余行，按原逻辑计算WAC
    for i in range(1, len(group)):
        prev_row = group.iloc[i-1]
        current_row = group.iloc[i]
        wac_col_idx = group.columns.get_loc('WAC')
        
        if current_row['Rec'] == 0 or current_row['Rec'] < 0:
            group.iloc[i, wac_col_idx] = prev_row['WAC']
        elif current_row['Rec'] > 0:
            group.iloc[i, wac_col_idx] = (current_row['Amt'] + prev_row['WAC'] * prev_row['Qtty']) / current_row['Qtty']
        else:
            group.iloc[i, wac_col_idx] = 0
    return group

步骤2：按Class分组并应用函数 #

使用groupby按Class分组，对每个分组调用上述函数，最后合并结果并保持原行顺序：

# 按Class分组计算，group_keys=False避免添加分组键索引
result_df = df.groupby('Class', group_keys=False).apply(calculate_wac)
# 重置索引以保持原DataFrame的行顺序
result_df = result_df.sort_index()

# 输出结果
print(result_df)

说明 #

• 该方案会为每个Class分组独立计算WAC，分组间的计算互不干扰，完全匹配原逻辑的分组需求。
• 计算后的结果会保留原DataFrame的行顺序，与目标输出一致。

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