Pandas groupby().agg() 中自定义加权平均计算的闭包应用

import pandas as pd
import numpy as np

# 原始数据
df_original = pd.DataFrame({
    'id': [1, 1, 2, 2, 3],
    'amount': [10, 200, 1, 10, 150],
    'other_col': [0.1, 0.6, 0.7, 0.2, 0.4]
})

# 尝试直接定义的加权平均函数 (会导致错误)
def problematic_weighted_mean(x):
    try:
        # 这里的 df_original.loc[x.index, 'amount'] 会导致 NameError
        # 因为当 agg 调用此函数时，df_original 不在其作用域内
        return np.average(x, weights=df_original.loc[x.index, 'amount']) > 0.5
    except ZeroDivisionError:
        return 0

# 尝试在另一个函数中调用聚合
def some_function_problematic(df1):
    # 这里的 df1 传入后，problematic_weighted_mean 函数并不能直接访问到它
    df1_result = df1.groupby('id').agg(
        xx=('amount', lambda x: x.sum() > 100),
        yy=('other_col', problematic_weighted_mean) # 这里的调用会出错
    ).reset_index()
    return df1_result

# df_result = some_function_problematic(df_original.copy()) # 运行会报错

import pandas as pd
import numpy as np

# 1. 定义一个创建加权平均聚合器的外部函数
def create_weighted_mean_aggregator(df_context):
    """
    创建一个闭包，用于计算指定列的加权平均。
    df_context: 原始DataFrame，用于提供权重列的上下文。
    """
    def inner_weighted_mean(x):
        """
        实际执行加权平均计算的内部函数。
        x: 当前分组中被聚合的Series (例如 'other_col')。
        """
        try:
            # 内部函数可以访问外部函数作用域中的 df_context
            # x.index 确保我们获取到当前分组对应的权重
            weights = df_context.loc[x.index, 'amount']
            # 避免所有权重都为0导致ZeroDivisionError
            if weights.sum() == 0:
                return 0
            return np.average(x, weights=weights) > 0.5
        except ZeroDivisionError:
            # 当所有权重都为0时，np.average 可能会抛出此错误
            return 0
    return inner_weighted_mean

# 2. 定义主函数，它将DataFrame作为参数并执行聚合
def process_data_with_weighted_mean(df_input=None):
    """
    处理DataFrame，计算分组总和和加权平均。
    df_input: 待处理的DataFrame。
    """
    if df_input is None:
        raise ValueError("Input DataFrame cannot be None.")

    # 3. 创建一个绑定了当前 df_input 上下文的加权平均函数实例
    # 这里 create_weighted_mean_aggregator(df_input) 返回的是 inner_weighted_mean 函数
    weighted_mean_for_current_df = create_weighted_mean_aggregator(df_input)

    # 4. 在 groupby().agg() 中使用这个闭包实例
    df_result = df_input.groupby('id').agg(
        xx=('amount', lambda x: x.sum() > 100),  # 示例：检查amount总和是否大于100
        yy=('other_col', weighted_mean_for_current_df) # 使用闭包函数进行加权平均
    ).reset_index()
    return df_result

# 示例数据
df_data = pd.DataFrame({
    'id': [1, 1, 2, 2, 3],
    'amount': [10, 200, 1, 10, 150],
    'other_col': [0.1, 0.6, 0.7, 0.2, 0.4]
})

# 调用主函数处理数据
df_processed = process_data_with_weighted_mean(df_data.copy())
print(df_processed)

   id     xx     yy
0   1   True   True
1   2  False  False
2   3   True  False