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