需求背景 #

你最初实现了针对特定列的分组加权均值计算，但希望把这个逻辑抽象成可复用的函数，支持不同的DataFrame、待计算列、权重列和分组列，避免重复写相似代码或lambda表达式，这个思路特别实用！

最终实现的可复用函数 #

你已经搞定了参数传递的问题，我把这个函数整理得更清晰，还补充了一些细节让它更健壮：

import numpy as np
import pandas as pd

def calc_weighted_avg(basket: pd.DataFrame, weight_col: str, used_col: list[str] | str, agg_col: str) -> pd.DataFrame:
    """
    按指定列分组，计算目标列的加权均值
    
    参数:
        basket: 输入的源DataFrame
        weight_col: 用作权重计算的列名称
        used_col: 需要计算加权均值的列（支持单列字符串或多列列表格式）
        agg_col: 用于分组聚合的列名称
    
    返回:
        包含各分组加权均值结果的DataFrame
    """
    def weighted_mean(x):
        # 先去除空值，避免后续计算出错
        y = x.copy().dropna()
        if y.empty:
            return np.nan
        # 从原DataFrame中匹配当前分组行对应的权重值
        weights = basket.loc[y.index, weight_col]
        return np.average(y, weights=weights)
    
    # 兼容单列输入的场景，统一转为列表格式处理
    if isinstance(used_col, str):
        used_col = [used_col]
    
    # 执行分组与聚合计算
    weight_avg_table = basket[used_col].groupby(basket[agg_col]).agg(weighted_mean)
    return weight_avg_table

使用示例 #

假设你有这样的示例数据：

trade_basket = pd.DataFrame({
    'Side': ['Buy', 'Sell', 'Buy', 'Sell', 'Buy'],
    'NotionalTraded': [100, 200, 150, 300, 250],
    'Price': [10.5, 10.7, 10.6, 10.8, 10.4],
    'Volume': [50, 60, 55, 65, 45]
})

你可以灵活调用这个函数：

# 按Side分组，用NotionalTraded做权重，计算Price和Volume的加权均值
result = calc_weighted_avg(
    basket=trade_basket,
    weight_col='NotionalTraded',
    used_col=['Price', 'Volume'],
    agg_col='Side'
)
print(result)

额外优化说明 #

1. 类型注解：添加了类型提示，让函数的参数和返回值更清晰，也方便IDE给出智能提示
2. 单列兼容：处理了used_col为单字符串的情况，不用每次都手动转成列表
3. 文档字符串：补充了docstring，后续自己或团队成员能快速理解函数的用途和参数含义

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