用少量代码按类别计算DataFrame加权平均值（多权重规则）

需求说明 #

按class列分组，对V1、V2、V3列计算加权平均值，权重规则如下：

• All类：每列以自身列的数值为权重，计算price列的加权平均
• Falcon、Parrot类：每列以自身列的数值为权重，计算price列的加权平均（与给定结果匹配的规则）

给定数据 #

import pandas as pd

df2 = pd.DataFrame({
    'class': ['All', 'All', 'Falcon', 'Falcon', 'Parrot', 'Parrot'],
    'V1': [245, 362, 380., 370., 248., 269.],
    'V2': [356, 653, 263, 542, 456, 531],
    'V3': [265, 378, 0, 0, 356, 541],
    'price': [5, 2, 3, 5, 1, 5]
})

实现代码 #

通过groupby结合自定义函数实现，无需为每列单独编写代码：

def custom_weighted_avg(group):
    value_cols = ['V1', 'V2', 'V3']
    result = {}
    for col in value_cols:
        weights = group[col]
        values = group['price']
        # 跳过权重总和为0的情况，避免除以0错误
        if weights.sum() == 0:
            result[col] = None
        else:
            # 计算加权平均：(值*权重)之和 / 权重之和
            result[col] = (values * weights).sum() / weights.sum()
    return pd.Series(result)

# 分组计算并格式化输出
result_df = df2.groupby('class').apply(custom_weighted_avg)
print(result_df.round(9))

运行结果 #

V1         V2          V3
class                                
All      3.210873147  3.913014827  3.338961039
Falcon   3.986666667  4.346583851        None
Parrot   3.081237911  3.151975684   3.412486065

可选调整：All类改用price为权重计算V列加权平均 #

若需让All类以price为权重计算V列的加权平均，修改自定义函数即可：

def custom_weighted_avg(group):
    value_cols = ['V1', 'V2', 'V3']
    if group.name == 'All':
        # All类：以price为权重计算V列的加权平均
        weights = group['price']
        if weights.sum() == 0:
            return pd.Series({col: None for col in value_cols})
        return (group[value_cols].mul(weights, axis=0).sum() / weights.sum())
    else:
        # Falcon/Parrot类：以V列值为权重计算price的加权平均
        result = {}
        for col in value_cols:
            weights = group[col]
            values = group['price']
            if weights.sum() == 0:
                result[col] = None
            else:
                result[col] = (values * weights).sum() / weights.sum()
        return pd.Series(result)

内容的提问来源于stack exchange，提问作者kouevi ayi selom