如何在Pandas中保留按ID1分组后ID2重复出现的行？

问题描述 #

原始DataFrame如下：

ID 1          ID 2
1             5
1             5
1             6
2             7
2             5
2             7
3             8
3             9
3             10

需求：按ID 1分组后，保留每个分组内ID 2出现次数大于1的所有行。

已尝试代码：

df_temp = df.groupby('ID 1')
df_output = df_temp['ID 2'].value_counts()[df_temp['ID 2'].value_counts() > 1]

得到的中间结果：

ID 1          ID 2
1             5         2
2             7         2

需要基于这个中间结果，筛选原始DataFrame得到目标结果：

ID 1          ID 2
1             5
1             5
2             7
2             7

解决方法 #

方法一：利用已生成的df_output筛选 #

df_output是多层索引的Series，先将其转换为包含ID 1和ID 2的筛选表，再匹配原始数据：

# 将中间结果转为筛选用DataFrame
filter_df = df_output.reset_index()[['ID 1', 'ID 2']]
# 内连接匹配原始数据
result = df.merge(filter_df, on=['ID 1', 'ID 2'], how='inner')

或者用元组匹配的方式：

# 获取需要保留的(ID1, ID2)组合
keep_pairs = df_output.index.to_list()
# 逐行判断是否在保留组合内
result = df[df.apply(lambda row: (row['ID 1'], row['ID 2']) in keep_pairs, axis=1)]

方法二：直接用transform一步完成（更高效） #

无需生成中间结果，直接在分组内计算ID 2的出现次数并筛选：

# 按ID1+ID2分组，计算每个组合的出现次数，筛选次数>1的行
result = df[df.groupby(['ID 1', 'ID 2'])['ID 2'].transform('size') > 1]

以上两种方法均可得到目标结果，其中方法二代码更简洁，执行效率也更高。

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