计算每组前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

所需的DF：

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值，创建前两行的子集数据帧，求和，然后用值填充字典。然后映射回字典-但我的数据帧非常大，这需要永远的时间。有什么想法吗？

谢谢!

Answer 1

您可以使用transform为每行获取一个新值和一个lambda函数。在lambda中，您可以使用head(2)获取每个组的前2行，然后对它们执行sum()操作：

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