如何用正则表达式从Pandas DataFrame列生成指定映射字典？

问题背景 #

现有如下Pandas DataFrame：

import pandas as pd
import numpy as np
df = pd.DataFrame({'Original': [92,93,94,95,100,101,102],
             'Sub_90': [99,98,99,100,102,101,np.nan],
             'Sub_80': [99,98,99,100,102,np.nan,np.nan],
             'Gen_90': [99,98,99,100,102,101,101],
             'Gen_80': [99,98,99,100,102,101,100]})

需要生成如下映射字典：

{
    'Gen_90': 'Original',
    'Sub_90': 'Gen_90',
    'Gen_80': 'Original',
    'Sub_80': 'Gen_80',
 }

实际数据还包含Gen_70、Gen_60……Gen_10以及Sub_70、Sub_60……Sub_10这类列，需通过正则实现两个规则：

• 所有Gen_xx列映射到Original；
• 所有Sub_xx列映射到同编号的Gen_xx列。

实现方法 #

可以通过正则匹配列名的前缀与编号，快速生成目标映射字典：

方法1：循环遍历+正则匹配 #

import re

# 获取DataFrame的所有列名
columns = df.columns.tolist()

mapping = {}
for col in columns:
    # 处理Gen_xx格式的列
    gen_match = re.match(r'Gen_(\d+)', col)
    if gen_match:
        mapping[col] = 'Original'
        continue
    # 处理Sub_xx格式的列
    sub_match = re.match(r'Sub_(\d+)', col)
    if sub_match:
        num = sub_match.group(1)
        mapping[col] = f'Gen_{num}'

print(mapping)

方法2：简洁字典推导式 #

import re

mapping = {
    col: 'Original' if re.match(r'Gen_\d+', col) 
    else f'Gen_{re.match(r"Sub_(\d+)", col).group(1)}' 
    for col in df.columns 
    if re.match(r'(Gen|Sub)_\d+', col)
}

print(mapping)

关键说明 #

• 正则r'Gen_(\d+)'匹配以Gen_开头后接数字的列，捕获组提取编号；
• 正则r'Sub_(\d+)'匹配以Sub_开头的列，提取编号后拼接成对应Gen_xx列名；
• 两种方法会自动覆盖所有符合格式的Gen_xx和Sub_xx列，无需手动指定编号。

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