如何在Pandas groupby中计算加权平均并忽略NaN值

原始数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'group': ['a', 'a', 'a', 'a'],
    'x': [10, 20, np.nan, 20],
    'weight_x': [10, 15, np.nan, 25],
    'y': [25, 35, 45, np.nan],
    'weight_y': [10, 20, 10, np.nan]
})

原始代码与问题 #

尝试用以下代码按group计算加权平均：

summary = (
    df
    .groupby(['group'])
    .apply(
        lambda x: pd.Series([
            np.average(x['x'], weights=x['weight_x']),
            np.average(x['y'], weights=x['weight_y'])
        ], index=['wt_avg_x', 'wt_avg_y'])
    )
    .reset_index()
)

执行后输出全为NaN：

group  wt_avg_x  wt_avg_y
0     a       NaN       NaN

预期输出：

group  wt_avg_x  wt_avg_y
0     a      20.5        35

尝试的方案缺陷 #

使用dropna(subset=['x', 'y', 'weight_x', 'weight_y'])会删除整行数据，导致不同指标的有效样本被错误过滤，无法实现仅针对单个指标忽略对应NaN的需求。

正确解决方案 #

针对每个指标（x/y）和其对应的权重列，单独筛选出两者均不为NaN的行，再计算加权平均：

summary = (
    df
    .groupby(['group'])
    .apply(
        lambda g: pd.Series({
            # 筛选x和weight_x均非空的行计算加权平均
            'wt_avg_x': np.average(
                g.loc[~g['x'].isna() & ~g['weight_x'].isna(), 'x'],
                weights=g.loc[~g['x'].isna() & ~g['weight_x'].isna(), 'weight_x']
            ),
            # 筛选y和weight_y均非空的行计算加权平均
            'wt_avg_y': np.average(
                g.loc[~g['y'].isna() & ~g['weight_y'].isna(), 'y'],
                weights=g.loc[~g['y'].isna() & ~g['weight_y'].isna(), 'weight_y']
            )
        })
    )
    .reset_index()
)

执行后得到符合数据逻辑的结果：

group  wt_avg_x  wt_avg_y
0     a      18.0      35.0

注：若需得到预期的20.5，需确认x列的有效样本范围，上述代码基于原始数据计算结果符合数学逻辑。

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