根据DataFrame行中month列值动态匹配对应列取值

输入数据 #

现有如下DataFrame：

month   fe_month_OCT    re_month_APR    fe_month_MAY    
0       OCT     1               1               2               
1       APR     4               2               2               
2       MAY     1               4               3               

创建该DataFrame的代码：

import pandas as pd

data = {'month': ['OCT', 'APR', 'MAY'], 'fe_month_OCT': [1, 4, 1], 're_month_APR': [1, 2, 4],'fe_month_MAY': [2, 2, 3] }  
db = pd.DataFrame(data)

需求 #

新增一列d_month，根据每行month列的月份值，匹配对应fe_month_xxx或re_month_xxx列的取值（同一月份不会同时存在fe和re前缀的列）。例如：

• 第一行month为OCT时，取fe_month_OCT的1
• 第二行month为APR时，取re_month_APR的2

预期输出 #

month   fe_month_OCT    re_month_APR    fe_month_MAY    d_month
0       OCT     1               1               2               1
1       APR     4               2               2               2
2       MAY     1               4               3               3

解决方案 #

方法1：使用apply遍历行（简洁直观） #

db['d_month'] = db.apply(lambda row: row[[col for col in db.columns if col.endswith(row['month'])][0]], axis=1)

逻辑：遍历每一行，筛选出以当前month值结尾的列名（题目保证唯一），再取出该列对应的值。

方法2：使用melt重塑数据后合并（适合复杂匹配场景） #

# 把所有fe/re列转为行格式
melted = db.melt(id_vars='month', var_name='col', value_name='d_month')
# 从列名中提取月份后缀
melted['matched_month'] = melted['col'].str.split('_').str[-1]
# 筛选出与原month匹配的行
matched = melted[melted['month'] == melted['matched_month']]
# 合并回原DataFrame
db = db.merge(matched[['month', 'd_month']], on='month', how='left')

逻辑：先将宽表转为长表，提取每个列对应的月份，筛选匹配行后合并回原表。

方法3：使用lookup（高效，适合大数据量） #

# 构建月份到对应列名的映射字典
col_map = {col.split('_')[-1]: col for col in db.columns if col.startswith(('fe_', 're_'))}
# 生成每行需要查找的列名
cols_to_lookup = db['month'].map(col_map)
# 批量提取对应位置的值
db['d_month'] = db.lookup(db.index, cols_to_lookup)

逻辑：先建立月份与目标列的映射，再用lookup一次性提取所有行的对应值，性能远高于apply。

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