pandas DataFrame使用for循环条件判断拆分列的错误排查咨询

原代码错误点 #

• 语法层面：混淆了赋值运算符=和比较运算符==，所有修改列的位置都用了==，只会返回布尔判断结果，不会真正修改DataFrame的数值
• 方法调用错误：str.contains()不支持直接传入列表作为匹配参数，需要先把待匹配的前缀拼接为正则或格式
• 逻辑错误：在iterrows循环中直接对整列赋值，就算语法正确也会覆盖全列数值，无法实现按行匹配的需求；同时直接把列表l赋值给col3也不符合提取对应前缀的要求
• pandas最佳实践：这种结构化的行匹配需求优先用向量化操作实现，比iterrows遍历效率高且不易出错

正确实现代码 #

import pandas as pd

# 初始化数据
l = ['AF_', 'LT_']
df = pd.DataFrame({'col1': [1, 2, 3, 4, 5], 'col2': ['AF_123_PX', 'OX', 'AF', 'BX', 'LT_123']})

# 构造匹配正则，^表示仅匹配开头，避免字符串中间出现对应前缀被误匹配
match_pattern = '|'.join(l)
# 提取匹配到的前缀
df['match_prefix'] = df['col2'].str.extract(f'^({match_pattern})')
# 生成匹配掩码
mask = df['match_prefix'].notna()

# 处理匹配成功的行
df.loc[mask, 'col3'] = df.loc[mask, 'match_prefix'].str.rstrip('_')
df.loc[mask, 'col4'] = df.loc[mask, 'col2'].str.split('_', n=1).str[1]

# 处理匹配失败的行
df.loc[~mask, 'col3'] = 'Global'
df.loc[~mask, 'col4'] = df.loc[~mask, 'col2']

# 删除临时辅助列
df = df.drop(columns=['match_prefix'])

运行上述代码后得到的DataFrame与预期结果完全一致。

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