在Pandas中分割聚合数据的更好方法是什么？

Question

我正在尝试对两列进行分组以获取聚合值，然后测试该值以查看它是大于还是小于阈值。我有什么：

SEGMENT = df.groupby(['Col_1','Col_2'])['Number'].apply(lambda x: '1_5' if sum(x) <6 else '6+'

很慢。这种方法有根本错误吗？感谢。

编辑：

SEGMENT = df.groupby(['Col_1','Col_2'])['Number'].sum().apply(lambda x: '1_5' if x <6 else '6+'

这加快了3倍。

Answer 1

您可以执行transform并将其用作布尔掩码：

g = df.groupby(['Col_1','Col_2'])

mask = g["Number"].transform("sum") < 6

df[mask]  # with group sum smaller than 6
df[~mask]  # with group sum greater or equal 6

您也可以使用filter：

g.filter(lambda x: x.sum() >= 6)