在数据分析和处理中，我们经常需要根据某些规则为dataframe中的数据添加分类标签。一个常见的场景是，我们有一个包含关键词到类别的映射字典，并希望根据dataframe某一列文本内容中是否包含这些关键词来为其分配相应的类别。然而，当字典的键并非dataframe列中的精确值，而是其子字符串时，标准的dataframe.map()方法便无法直接满足需求。本文将深入探讨如何优雅地解决这一问题。
遇到的挑战
假设我们有一个商品列表DataFrame，其中包含商品名称（Item列），以及一个将商品关键词映射到其类别的字典。例如：
category_dict = {
'apple': 'fruit',
'grape': 'fruit',
'chickpea': 'beans',
'coffee cup': 'tableware'
}
data = {
'Item': [
'apple from happy orchard',
'grape from random vineyard',
'chickpea and black bean mix',
'coffee cup with dog decal'
],
'Cost': [15, 20, 10, 14]
}
df = pd.DataFrame(data)我们期望的结果是为DataFrame添加一个Category列，根据Item列中的关键词从category_dict中查找对应的类别。如果直接使用df['Item'].map(category_dict)，由于Item列中的值（如"apple from happy orchard"）与字典键（"apple"）不完全匹配，map方法将返回NaN，无法达到预期效果。
解决方案：结合apply与自定义lambda函数
解决此问题的核心在于对DataFrame的每一行（或具体到每一单元格）应用一个自定义逻辑，该逻辑能够遍历字典，检查字典键是否为单元格文本的子字符串。Pandas的apply()方法结合Python的lambda函数和生成器表达式，能够高效地实现这一目标。
以下是实现这一功能的代码示例：
import pandas as pd
# 定义分类字典
category_dict = {
'apple': 'fruit',
'grape': 'fruit',
'chickpea': 'beans',
'coffee cup': 'tableware'
}
# 创建示例DataFrame
data = {
'Item': [
'apple from happy orchard',
'grape from random vineyard',
'chickpea and black bean mix',
'coffee cup with dog decal',
'banana smoothie' # 添加一个没有匹配项的示例
],
'Cost': [15, 20, 10, 14, 12]
}
df = pd.DataFrame(data)
# 使用apply和lambda函数添加'Category'列
df['Category'] = df['Item'].apply(
lambda item_text: next(
(value for key, value in category_dict.items() if key in item_text),
None
)
)
print("原始DataFrame：")
print(pd.DataFrame(data))
print("\n添加分类列后的DataFrame：")
print(df)运行上述代码，将得到如下输出：
原始DataFrame：
Item Cost
0 apple from happy orchard 15
1 grape from random vineyard 20
2 chickpea and black bean mix 10
3 coffee cup with dog decal 14
4 banana smoothie 12
添加分类列后的DataFrame：
Item Cost Category
0 apple from happy orchard 15 fruit
1 grape from random vineyard 20 fruit
2 chickpea and black bean mix 10 beans
3 coffee cup with dog decal 14 tableware
4 banana smoothie 12 None代码解析
df['Item'].apply(...): apply()方法是Pandas DataFrame或Series的一个强大功能，它允许我们对Series中的每一个元素或DataFrame的每一行/列应用一个函数。在这里，我们将其应用于Item列，意味着对Item列中的每一个字符串执行一次指定的lambda函数。
lambda item_text:: 这是一个匿名函数，它接收一个参数item_text，代表Item列中的当前字符串（例如，"apple from happy orchard"）。
-
next((value for key, value in category_dict.items() if key in item_text), None): 这是解决方案的核心逻辑。
- (value for key, value in category_dict.items() if key in item_text): 这是一个生成器表达式。它遍历category_dict中的每一个键值对（key, value）。对于每个键key，它检查key是否作为子字符串存在于当前的item_text中。如果条件key in item_text为真，则生成器会产生对应的value。
-
next(generator, default): next()函数用于从迭代器（这里是生成器表达式）中获取下一个元素。
- 当生成器产生第一个匹配的value时，next()会立即返回这个value。这意味着一旦找到第一个匹配的关键词，就会停止搜索，并返回该关键词对应的类别。
- 如果生成器表达式遍历完整个字典，但没有找到任何匹配的key，next()函数会返回其第二个参数None。这确保了即使没有找到匹配项，Category列也不会引发错误，而是填充None。
注意事项与最佳实践
匹配顺序的重要性：next()函数会返回第一个匹配到的项。如果字典中存在重叠的关键词（例如，{'apple': 'fruit', 'red apple': 'red fruit'}），且item_text中同时包含这两个关键词，那么返回的类别将取决于category_dict.items()的遍历顺序。在Python 3.7+中，字典会保留插入顺序，因此通常会返回先插入的那个匹配项。如果匹配顺序很重要，请确保字典的定义顺序或考虑更复杂的匹配逻辑（例如，优先匹配更长的关键词）。
性能考量：对于非常大的DataFrame，apply()方法虽然功能强大，但在纯Python循环中执行自定义逻辑可能会比Pandas内置的向量化操作慢。对于数百万行的数据，可以考虑使用更底层的字符串匹配库或对字典进行预处理（例如，使用正则表达式）来优化性能。然而，对于大多数常见规模的数据集，apply()方法通常足够高效。
大小写敏感性：key in item_text是大小写敏感的。如果需要进行不区分大小写的匹配，应在比较前将key和item_text都转换为小写或大写，例如key.lower() in item_text.lower()。
处理无匹配项：next()函数中的None作为默认值是处理无匹配项的优雅方式。根据业务需求，也可以将其替换为其他默认值（如'Other'、'Unknown'等），或者进行后续处理来识别未分类的数据。
总结
通过结合使用Pandas的apply()方法、Python的lambda函数和生成器表达式，我们可以灵活高效地为DataFrame添加基于子字符串匹配的分类列。这种方法不仅解决了传统map()方法的局限性，还通过next()函数的默认值参数，优雅地处理了无匹配项的情况。理解并掌握这一技巧，将极大地提升您在处理复杂文本数据分类任务时的效率和代码质量。
