pandas按多列分组并应用apply函数生成新列

Question

我可以使用pandas的groupby-apply功能在数据框中生成新的列。

例如，我有一个数据框：

df = pd.DataFrame({'A':[1,2,3,4],'B':['A','B','A','B'],'C':[0,0,1,1]})

请尝试通过 groupby-apply 生成一个名为'D'的新列。

这个方法是可行的：

df = df.assign(D=df.groupby('B').C.apply(lambda x: x - x.mean()))

我认为它返回一个具有与数据框相同索引的系列：

In [4]: df.groupby('B').C.apply(lambda x: x - x.mean())
Out[4]:
0   -0.5
1   -0.5
2    0.5
3    0.5
Name: C, dtype: float64

但是如果我尝试使用多个列生成新的列，我不能直接将其分配给新的列。因此，这不起作用：

 df.assign(D=df.groupby('B').apply(lambda x: x.A - x.C.mean()))

返回

TypeError: incompatible index of inserted column with frame index

事实上，groupby-apply返回：

In [8]: df.groupby('B').apply(lambda x: x.A - x.C.mean())
Out[8]:
B
A  0    0.5
   2    2.5
B  1    1.5
   3    3.5
Name: A, dtype: float64

我可以做到。

df.groupby('B').apply(lambda x: x.A - x.C.mean()).reset_index(level=0,drop=True))

但是它看起来很啰嗦，我不确定它是否总是按预期工作。

所以我的问题是：（i）当使用pandas groupby-apply时，何时返回类似索引的系列和多级索引系列？ （ii）有没有更好的方法将新列分配给groupby-apply并应用于多个列？

- Jongmmm

Answer 1

对于这种情况，我认为在apply中不需要包含列A，我们可以使用transform。

df.A-df.groupby('B').C.transform('mean')
Out[272]: 
0    0.5
1    1.5
2    2.5
3    3.5
dtype: float64

你可以将其重新分配回去

df['diff']= df.A-df.groupby('B').C.transform('mean')
df
Out[274]: 
   A  B  C  diff
0  1  A  0   0.5
1  2  B  0   1.5
2  3  A  1   2.5
3  4  B  1   3.5

- BENY

Answer 2

让我们在groupby中使用group_keys=False

df.assign(D=df.groupby('B', group_keys=False).apply(lambda x: x.A - x.C.mean()))

输出：

   A  B  C    D
0  1  A  0  0.5
1  2  B  0  1.5
2  3  A  1  2.5
3  4  B  1  3.5

- Scott Boston

---

这直接回答了我的问题。 - Jongmmm

请注意，此解决方案仅在.apply()方法返回与您要分配的数据框具有相同行数的系列时才有效。否则，返回的系列的索引和数据框的索引将不匹配。 - petrovski

@petrovski 感谢您指出这一点，我一直在想这样的事情可能如何运作。 - geo909