如何修改字典处理代码以实现DataFrame按ID清洗文本末尾指定数字

原字典处理代码 #

原代码用于处理字典中各DataFrame的列名后缀：

def remove(dic):
    suffix_map = {
        "id1": r'(?:0)$',
        "id2": r'(?:1)$',
        "id3": r'(?:2)$',
        "id4": r'(?:3)$',
    }

    for key, val in dic.items():
        suffix_pattern = suffix_map.get(key)
        if suffix_pattern:
            val.columns = val.columns.str.replace(suffix_pattern, '', regex=True)
            val.columns = val.columns.str.replace(r'_$', '', regex=True)

    return dic

初始DataFrame #

需要处理的包含id和text列的DataFrame：

import pandas as pd

df1 = pd.DataFrame([{'id': "id1", 'text': 'x11'}, {'id': "id2", 'text': 'x342'}, {'id': "id2",'text': 'b34'}])

期望清洗结果 #

处理后需得到的DataFrame：

out = pd.DataFrame([{'id': "id1", 'text': 'x1'}, {'id': "id2", 'text': 'x34'}, {'id': "id2",'text': 'b34'}])

修改后的适配代码 #

针对单个DataFrame的text列处理，修改后的函数如下（已根据期望结果调整suffix_map的匹配规则）：

import pandas as pd
import re

def clean_text_by_id(df):
    suffix_map = {
        "id1": r'(?:1)$',  # 移除id1对应text末尾的1
        "id2": r'(?:2)$',  # 移除id2对应text末尾的2
        "id3": r'(?:3)$',
        "id4": r'(?:0)$',
    }
    
    # 按行匹配id对应的规则，清洗text列
    df['text'] = df.apply(
        lambda row: re.sub(suffix_map[row['id']], '', row['text'])
        if row['id'] in suffix_map else row['text'],
        axis=1
    )
    return df

# 调用示例
cleaned_df = clean_text_by_id(df1.copy())
print(cleaned_df)

代码说明 #

• 调整处理对象：从字典中DataFrame的列名，改为单个DataFrame的text列内容
• 通过apply遍历每行，根据当前行的id值匹配对应的正则规则
• 仅当id存在匹配规则时执行替换，无规则则保留原文本
• 后续新增id规则只需在suffix_map中添加对应键值对即可

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