Pandas Dataframe groupby 快速计算行数

Question

我有一个看起来像这样的数据框

Class_ID  Student_ID  feature
1         4           31
1         4           86
1         4           2
1         2           11
1         2           0
5         3           2
5         9           3
5         9           2

并且我想计算学生出现在每个

Class_ID

中的次数，所以期望的结果如下所示：

Class_ID  Student_ID  feature  count
1         4           31       3
1         4           86       3
1         4           2        3
1         2           11       2
1         2           0        2
5         3           2        1
5         9           3        2
5         9           2        2

我是这样做的：

df['dummy'] = 1
df['count'] = df.groupby(['Class_ID', 'Student_ID'], group_keys=False)['dummy'].transform(lambda x: x.sum())

它工作正常，但我的实际数据框相当大（〜1M行），并且代码非常慢，所以我想问是否有更快的方法/更好的方法来做到这一点？谢谢。

Answer 1

您可以避免创建虚拟列并用

transform

代替

"size"

：

df["count"] = df.groupby(["Class_ID", "Student_ID"]).transform("size")

在我的机器上，处理 100 万行时速度大约快 3/4 倍。

输出：

   Class_ID  Student_ID  feature  dummy  count
0         1           4       31      1      3
1         1           4       86      1      3
2         1           4        2      1      3
3         1           2       11      1      2
4         1           2        0      1      2
5         5           3        2      1      1
6         5           9        3      1      2
7         5           9        2      1      2