首页 > 解决方案 > 使用 pandas 中的分组 .agg 计算加权平均值

问题描述

我想使用pandas.agg()中的函数按组计算数据集中一列的平均值和另一列的加权平均值。我知道一些解决方案,但它们不是很简洁。

一个解决方案已在此处发布(pandas and groupby: how to calculate weighted averages within an agg,但它似乎仍然不是很灵活,因为 weights 列在 lambda 函数定义中是硬编码的。我正在寻找创建一个语法更接近于此:

(
df
.groupby(['group'])
.agg(avg_x=('x', 'mean'),
     wt_avg_y=('y', 'weighted_mean', weights='weight')
)

这是一个完整的示例,其中的代码似乎不必要地复杂:

import pandas as pd
import numpy as np

# sample dataset
df = pd.DataFrame({
    'group': ['a', 'a', 'b', 'b'],
    'x': [1, 2, 3, 4],
    'y': [5, 6, 7, 8],
    'weights': [0.75, 0.25, 0.75, 0.25]
})
df
#>>>    group   x   y   weights
#>>> 0      a   1   5   0.75
#>>> 1      a   2   6   0.25
#>>> 2      b   3   7   0.75
#>>> 3      b   4   8   0.25

# aggregation logic
summary = pd.concat(
    [
        df.groupby(['group']).x.mean(),
        df.groupby(['group']).apply(lambda x: np.average(x['y'], weights=x['weights']))
    ], axis=1
)
# manipulation to format the output of the aggregation
summary = summary.reset_index().rename(columns={'x': 'avg_x', 0: 'wt_avg_y'})

# final output
summary
#>>>    group   avg_x   wt_avg_y
#>>> 0      a   1.50    5.25
#>>> 1      b   3.50    7.25

标签: python-3.xpandaspandas-groupby

解决方案


在整个 DataFrame 上使用该.apply()方法是我能得到的最简单的解决方案,它不会对函数定义中的列名进行硬编码。

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'group': ['a', 'a', 'b', 'b'],
    'x': [1, 2, 3, 4],
    'y': [5, 6, 7, 8],
    'weights': [0.75, 0.25, 0.75, 0.25]
})

summary = (
    df
    .groupby(['group'])
    .apply(
        lambda x: pd.Series([
            np.mean(x['x']),
            np.average(x['y'], weights=x['weights'])
        ], index=['avg_x', 'wt_avg_y'])
    )
    .reset_index()
)
# final output
summary
#>>>    group   avg_x   wt_avg_y
#>>> 0      a   1.50    5.25
#>>> 1      b   3.50    7.25

推荐阅读