Python中如何用DataFrame的apply替换嵌套for循环

初始代码与需求 #

给定以下DataFrame和列表：

import pandas as pd

my_df = pd.DataFrame({'fruits': ['apple', 'banana', 'cherry', 'durian'], 
                      'check': [False, False, False, False]})
my_list = ['pp', 'ana', 'ra', 'cj', 'up', 'down', 'pri']

需求：判断fruits列的每个元素是否包含my_list中的任意子串，若包含则将对应check列设为True，期望结果如下：

fruits  check
0   apple   True
1  banana   True
2  cherry  False
3  durian  False

报错情况 #

尝试用apply实现时写出如下代码：

my_df['fruits'].apply(lambda x: True for i in my_list if i in x)

触发错误：TypeError: 'generator' object is not callable

解决方法 #

方法1：修正apply的lambda写法

报错原因是直接将生成器表达式作为lambda的返回值，需要用any()来判断是否存在匹配的子串——any()会遍历生成器并返回布尔值，判断是否有任意子串匹配：

my_df['check'] = my_df['fruits'].apply(lambda x: any(sub in x for sub in my_list))

该代码对每个水果字符串x，检查my_list中是否存在任意子串sub包含在x内，匹配则返回True，否则返回False，直接赋值给check列即可。

方法2：更高效的矢量化操作（推荐）

Pandas的str.contains()支持正则表达式，可将my_list中的子串用|拼接成正则模式，实现矢量化匹配，效率远高于apply循环（尤其适用于大数据集）：

pattern = '|'.join(my_list)
my_df['check'] = my_df['fruits'].str.contains(pattern)

这里'|'.join(my_list)生成正则模式'pp|ana|ra|cj|up|down|pri'，str.contains()会检查每个水果字符串是否匹配该模式（即包含任意一个目标子串）。

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