Pandas按指定列分组并筛选包含特定值的分组

原始数据 #

A                B
MEDUI5945189     GATE IN
MEDUI5945189     RAIL LOAD
MEDUI5945189     GATE OUT
EBKG04830245     LOADED ON VESSEL
EBKG04830245     GATE OUT
COAU7242812270   VESSEL DEPARTURE
COAU7242812270   GATE IN
COAU7242812270   CHANGE IN SHIPMENT ETA
COAU7242812270   GATE OUT FULL
EBKG04830245     CHANGE IN SHIPMENT ETA
EBKG04830245     RAIL UNLOAD
EBKG04830245     VESSEL DEPARTURE

需求 #

按列A分组，仅保留那些分组内至少有一行的列B包含RAIL字样的完整分组，预期结果如下：

A                B
MEDUI5945189     GATE IN
MEDUI5945189     RAIL LOAD
MEDUI5945189     GATE OUT
EBKG04830245     LOADED ON VESSEL
EBKG04830245     GATE OUT
EBKG04830245     CHANGE IN SHIPMENT ETA
EBKG04830245     RAIL UNLOAD
EBKG04830245     VESSEL DEPARTURE

解决方案 #

你可以通过两种方式实现这个筛选逻辑：

方法一：用transform直接生成筛选标记 #

import pandas as pd

# 假设数据已加载到DataFrame df中
# 给每行标记其所在分组是否包含含RAIL的行
has_rail = df.groupby('A')['B'].transform(lambda x: x.str.contains('RAIL').any())

# 筛选符合条件的完整分组
df_sel = df[has_rail]

逻辑说明：transform会把分组的判断结果（是否含RAIL）广播到该分组的每一行，生成和原DataFrame行数一致的布尔序列，最后用这个序列筛选即可保留完整的目标分组。

方法二：先提取有效分组再筛选 #

import pandas as pd

# 先找出所有包含RAIL的行对应的A列值
valid_group_ids = df[df['B'].str.contains('RAIL')]['A'].unique()

# 筛选这些分组的所有行
df_sel = df[df['A'].isin(valid_group_ids)]

逻辑说明：先定位所有含RAIL的行，提取它们的分组ID，再用这些ID筛选原DataFrame，同样能得到所有符合条件的完整分组。

内容的提问来源于stack exchange，提问作者just_learning