python - 计算每组前n行的总和

Question

我想做的是按 A 列分组，然后取前两行的总和，然后将该值分配为新列。示例如下:

东风:

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 值创建字典，创建前两行的子集数据框，求和，然后用值填充字典来尝试此操作。然后将字典映射回去——但我的数据框非常大，需要很长时间。有什么想法吗？

谢谢!

Answer 1

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

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