如何在Pandas的np.where语句中用字典实现列映射

问题背景 #

给定以下DataFrame：

import pandas as pd
foo = pd.DataFrame({'id': [1,1,1,2,2,2],
                    'time': [1,2,3,1,2,3],
             'col_id': ['ffp','ffp','ffp', 'hie', 'hie', 'ttt'],
             'col_a': [1,2,3,4,5,6],
             'col_b': [-1,-2,-3,-4,-5,-6],
                'col_c': [10,20,30,40,50,60]})

需求是创建新列col，根据col_id的取值，分别从col_a、col_b或col_c中提取对应值。原实现采用嵌套np.where，但当列数较多时代码冗余，希望通过以下映射字典简化实现：

dict_cols_matching = {"ffp" : "col_a", "hie": "col_b", "ttt": "col_c"}

实现方案 #

方案1：使用lookup方法（简洁直观） #

lookup可直接根据行索引和目标列名的匹配关系提取值，结合字典映射快速实现：

foo['col'] = foo.lookup(foo.index, foo['col_id'].map(dict_cols_matching))

注意：lookup在pandas 1.2.0及以上版本被标记为弃用，但当前仍可正常使用；若需长期兼容性，建议选择其他方案。

方案2：apply逐行映射（逻辑清晰） #

通过apply遍历每行，根据col_id从字典中获取目标列名，再提取对应值：

foo['col'] = foo.apply(lambda row: row[dict_cols_matching[row['col_id']]], axis=1)

该方法逻辑易懂，适合映射规则复杂的场景，但在大数据集下效率略低于向量化方法。

方案3：向量化实现（性能最优） #

利用numpy的索引机制实现完全向量化操作，处理大规模数据时效率最高：

# 建立列名到DataFrame列索引的映射
col_position = {col: idx for idx, col in enumerate(foo.columns)}
# 将col_id转换为目标列的索引位置
target_positions = foo['col_id'].map(dict_cols_matching).map(col_position)
# 提取对应位置的值
foo['col'] = foo.values[foo.index, target_positions]

方案4：melt+合并（扩展性强） #

如果后续需要频繁调整映射规则，可先将数据转为长格式再匹配：

# 将需要映射的列转为长格式
melted_df = foo.melt(
    id_vars=['id', 'time', 'col_id'],
    value_vars=['col_a', 'col_b', 'col_c'],
    var_name='source_col',
    value_name='col'
)
# 生成反向映射（列名对应col_id值）
reverse_map = {v: k for k, v in dict_cols_matching.items()}
melted_df['match_id'] = melted_df['source_col'].map(reverse_map)
# 匹配后合并回原数据
foo = foo.merge(
    melted_df[melted_df['col_id'] == melted_df['match_id']][['id', 'time', 'col']],
    on=['id', 'time'],
    how='left'
)

该方案代码稍长，但扩展性好，便于后续修改映射关系。

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