如何筛选pandas DataFrame中col A或col B值在指定列表内的行

性能说明 #

你原来用iterrows的实现方案完全不是最优选择，时间复杂度为O(n*m)（n为DataFrame行数、m为list1长度），大数据量下执行效率极低，而且会出现一行同时满足两个条件时重复输出的问题。
pandas的向量化操作底层为C实现，性能远高于Python层的循环遍历，优先选择内置方法实现需求。

实现代码 #

场景1：完全匹配（col A/col B的值等于list1中的某一项） #

首先要保证col A、col B的类型和list1元素类型一致，你给出的list1元素是字符串类型，如果你的DataFrame中两列是整数类型，需要先做类型对齐：

# 把list1转成整数类型，和df的列类型对齐
list1_int = [int(x) for x in list1]
# 筛选满足条件的行
filtered_df = df[(df['col A'].isin(list1_int)) | (df['col B'].isin(list1_int))]

如果两列本身就是字符串类型，直接用原list1即可：

filtered_df = df[(df['col A'].isin(list1)) | (df['col B'].isin(list1))]

场景2：子串匹配（col A/col B的字符串包含list1中的某一项作为子串） #

如果你原来代码里的in是要判断子串包含关系，用如下写法：

# 构造正则匹配规则
match_rule = '|'.join(list1)
mask = df['col A'].str.contains(match_rule, na=False) | df['col B'].str.contains(match_rule, na=False)
filtered_df = df[mask]

结果输出 #

直接打印筛选后的DataFrame即可：

print(filtered_df)

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