按组计算前n行的和

Question

我希望做的是按A列分组，然后取前两行的和，将该值分配为新列。以下是示例：

DF:

ColA   ColB
 AA      2
 AA      1
 AA      5
 AA      3
 BB      9
 BB      3
 BB      2
 BB      12
 CC      0
 CC      10
 CC      5
 CC      3

期望的数据框：

ColA   ColB    NewCol
 AA      2       3
 AA      1       3
 AA      5       3 
 AA      3       3
 BB      9       12
 BB      3       12
 BB      2       12
 BB      12      12
 CC      0       10
 CC      10      10 
 CC      5       10
 CC      3       10

对于AA来说，它查看ColB并取前两行的总和，将该总和分配给newCol。我尝试通过循环遍历唯一的ColA值，创建前两行的子数据帧进行求和，然后用值填充字典来创建字典。然后映射回数据帧 - 但我的数据帧非常大，这需要很长时间。有什么想法吗？谢谢！

- dmd7

Answer 1

你可以使用 transform 和 lambda 函数按每行获取新值。在 lambda 中，你可以使用 head(2) 获取每个组的前两行并将它们相加：sum()。

df.groupby('ColA')['ColB'].transform(lambda x: x.head(2).sum())

- Mykola Zotko