Pandas 根据 groupby 平均值生成新列

问题描述：

我有一个 DataFrame

>>> df = pd.DataFrame({'a':[1,1,1,2,2,2],
...                    'b':[10,20,20,10,20,20],
...                    'result':[100,200,300,400,500,600]})
... 
>>> df
   a   b  result
0  1  10     100
1  1  20     200
2  1  20     300
3  2  10     400
4  2  20     500
5  2  20     600

并希望创建一个新列，该列是“a”和“b”对应值的平均结果。我可以使用 groupby 获取这些值：

>>> df.groupby(['a','b'])['result'].mean()
a  b 
1  10    100
   20    250
2  10    400
   20    550
Name: result, dtype: int64

但不知道如何将其转换为原始 DataFrame 中的新列。最终结果应如下所示，

>>> df
   a   b  result  avg_result
0  1  10     100         100
1  1  20     200         250
2  1  20     300         250
3  2  10     400         400
4  2  20     500         550
5  2  20     600         550

我可以通过循环“a”和“b”的组合来实现这一点，但对于较大的数据集来说，这会变得非常缓慢和难以处理。可能有一种更简单、更快捷的方法。

解决方案 1：

你需要transform：

df['avg_result'] = df.groupby(['a', 'b'])['result'].transform('mean')

这将为您生成一个正确索引的 groupby 值列：

   a   b  result  avg_result
0  1  10     100         100
1  1  20     200         250
2  1  20     300         250
3  2  10     400         400
4  2  20     500         550
5  2  20     600         550

解决方案 2：

由于之前的答案（https://stackoverflow.com/a/33445035/6504287）是基于 pandas 的，因此我添加了基于 pyspark 的解决方案，如下所示：因此最好使用Window以下代码片段示例中的函数：

    windowSpecAgg  = Window.partitionBy('a', 'b')
    ext_data_df.withColumn('avg_result', avg('result').over(windowSpecAgg)).show()

上述代码是根据之前提供的解决方案（https://stackoverflow.com/a/33445035/6504287）中的示例编写的。

解决方案 3：

您需要重置索引，例如：

df.reset_index()

输出应该像你想要的那样