Pandas分组后过滤组并保留分组结构的实现方法

嘿，这个场景我太熟悉了——用groupby.filter确实会把分组结构打散成普通DataFrame，完全懂你想要保留分组信息又不想多遍历一次数据集的需求！咱们来拆解一下：

先确认你的效率判断是对的！ #

你说的没错：如果先筛选A>0的行再做groupby，确实会遍历DataFrame两次（一次筛选、一次分组）；而先分组再过滤组，只需要一次分组遍历，之后只需要筛选组名，计算量确实更低，大数据集下这个差异会很明显。

怎么保留分组结构？ #

直接用groupby对象的组信息生成字典是最直接的方式，而且只需要一次分组操作：

方法1：遍历分组时直接筛选 #

# 一次groupby，同时筛选符合条件的组并生成字典
grouped_dict = {name: group for name, group in df.groupby('A') if name > 0}

这个字典里的key就是符合条件的组名（A列的正值），value就是对应组的完整DataFrame，和GroupBy对象里的每个组结构完全一致，你可以直接通过grouped_dict[某个组名]访问对应组的数据。

方法2：先创建GroupBy对象再筛选组 #

如果之后还要用到原GroupBy对象的其他方法，也可以先把GroupBy对象存下来，再基于它的组名筛选：

# 先创建GroupBy对象，只做一次分组
gb = df.groupby('A')
# 筛选出符合条件的组名，再生成字典
valid_group_names = [name for name in gb.groups if name > 0]
grouped_dict = {name: gb.get_group(name) for name in valid_group_names}

这个方式和第一种效率差不多，好处是可以复用gb对象做其他操作。

为什么groupby.filter不保留分组？ #

其实filter方法的设计初衷就是把符合条件的所有组的行合并成一个扁平的DataFrame，它的目标是输出过滤后的行数据，而不是保留分组结构，所以确实会丢掉分组信息。如果需要保留分组，手动提取符合条件的组是目前最稳妥的方式。

内容的提问来源于stack exchange，提问作者Pěťa Poliak