如何使用Pandas筛选包含Model A及其他Model的Folder数据？

需求说明 #

筛选出所有满足「Folder中既包含Model A，又包含B、C或D中任一Model」的行。

示例数据 #

import pandas as pd

data = {
    'Folder': [123, 123, 123, 4541, 4541, 4541, 4541, 11, 11, 222, 222, 222, 222],
    'Model': ['A', 'A', 'A', 'A', 'B', 'C', 'A', 'B', 'C', 'D', 'D', 'B', 'A']
}
df = pd.DataFrame(data)

解决方案一：使用groupby.filter() #

直接对分组应用过滤条件，一步得到结果：

def filter_group(group):
    # 检查分组是否包含A，且包含B/C/D中的至少一个
    has_a = 'A' in group['Model'].values
    has_non_a = group['Model'].isin(['B', 'C', 'D']).any()
    return has_a and has_non_a

# 应用过滤，返回符合条件的所有行
result = df.groupby('Folder').filter(filter_group)

解决方案二：先标记有效Folder再筛选 #

先计算每个Folder是否符合条件，再筛选原始数据：

# 计算每个Folder的两个条件标记
folder_check = df.groupby('Folder')['Model'].agg(
    has_a=lambda x: 'A' in x,
    has_non_a=lambda x: x.isin(['B', 'C', 'D']).any()
)

# 筛选同时满足两个条件的Folder
valid_folders = folder_check[(folder_check['has_a'] & folder_check['has_non_a'])].index

# 从原始DataFrame中提取有效Folder的所有行
result = df[df['Folder'].isin(valid_folders)]

最终结果 #

执行上述代码后，result的输出如下：

Folder Model
3     4541     A
4     4541     B
5     4541     C
6     4541     A
9      222     D
10     222     D
11     222     B
12     222     A

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