Pandas Group 多列加权平均

Question

假设我有以下数据框：

>>> df=pd.DataFrame({'category':['a','a','b','b'],
... 'var1':np.random.randint(0,100,4),
... 'var2':np.random.randint(0,100,4),
... 'weights':np.random.randint(0,10,4)})
>>> df
  category  var1  var2  weights
0        a    37    36        7
1        a    47    20        1
2        b    33     7        6
3        b    16     6        8

我可以这样计算“var1”的加权平均值：

>>> Grouped=df.groupby('category')
>>> GetWeightAvg=lambda g: np.average(g['var1'], weights=g['weights'])
>>> Grouped.apply(GetWeightAvg)
category
a    38.250000
b    23.285714
dtype: float64

但是我想知道是否有一种方法可以编写我的函数并将其应用于我的分组对象，以便我可以在应用它时指定我想要计算的列（或两者）。我希望能够在应用函数时指定，而不是将“var1”写入我的函数中。

就像我可以获得两列的未加权平均值一样：

>>> Grouped[['var1','var2']].mean()
          var1  var2
category            
a         42.0  28.0
b         24.5   6.5

我想知道是否有一种并行的方法可以使用加权平均值来做到这一点。

Answer 1

您可以应用并返回两个平均值：

In [11]: g.apply(lambda x: pd.Series(np.average(x[["var1", "var2"]], weights=x["weights"], axis=0), ["var1", "var2"]))
Out[11]:
               var1       var2
category
a         38.250000  34.000000
b         23.285714   6.428571

你可以把这个稍微简洁一点的写成一个函数：

In [21]: def weighted(x, cols, w="weights"):
             return pd.Series(np.average(x[cols], weights=x[w], axis=0), cols)

In [22]: g.apply(weighted, ["var1", "var2"])
Out[22]:
               var1       var2
category
a         38.250000  34.000000
b         23.285714   6.428571