Pandas通过分组计算平均值创建新列

Question

我有一个 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'的组合来实现，但是对于更大的数据集，这样做会变得非常缓慢和笨重。可能有一种更简单、更快捷的方法。

- rurp

---

我想强调这个问题的写作非常好，并且它的最小示例代码使其对未来的读者非常有用。 - mafu

Answer 1

你需要使用transform：

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

这将为您生成正确索引的分组值列：

   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

- Alex Riley

---

现在需要重新审视这个问题，因为有一些更新，上述代码返回错误 - TypeError: 'GroupedData' object is not subscriptable。 - Mithun Theertha

@MithunTheertha：这段代码在1.5.3版本上对我来说运行良好。你是在pyspark下运行吗？pandas没有GroupedData对象，但是pyspark有。 - Alex Riley

是的，我同意，后来我想通了，谢谢。 - Mithun Theertha

Answer 2

由于之前的答案(https://dev59.com/O1wX5IYBdhLWcg3wvRkf#33445035)是基于pandas的，我在下面添加了基于pyspark的解决方案： 因此最好使用Window函数，如下面的代码片段示例：

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

以上代码是针对之前提供的解决方案中所使用的示例（https://dev59.com/O1wX5IYBdhLWcg3wvRkf#33445035）而编写的。

- Mithun Theertha

Answer 3

你需要重置索引，像这样：

df.reset_index()

输出应该是你想要的那样。

- Muhammad Adli

---

你的答案可以通过提供更多支持性信息来改进。请编辑以添加更多细节，例如引用或文档，以便其他人能够确认你的答案是否正确。你可以在帮助中心找到关于如何撰写好答案的更多信息。 - Community