如何使用类IF语句根据条件从多列中取值创建新列

解决方案 #

首先澄清一个误解：np.where 并不是只能赋值固定值，它的第二、第三参数支持传入Series对象，可直接取对应行的列值，二分类场景可直接使用。

第一步：构造示例数据

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'field_a': [1,1,3,3],
    'field_b': [2,2,4,4],
    'which_field': ['a','b','a','b']
})

方案1：二分类简单场景，直接用np.where

df['output'] = np.where(df['which_field'] == 'a', df['field_a'], df['field_b'])

方案2：多条件映射场景（支持N个取值对应N个列，性能最优）

如果which_field有多个不同取值对应不同列，只需维护映射规则即可，无需修改核心逻辑：

# 定义映射关系：which_field取值 -> 对应的源列名，可无限新增规则
field_map = {
    'a': 'field_a',
    'b': 'field_b'
    # 'c': 'field_c',
    # 'd': 'field_d'
}

# 生成每行需要读取的目标列名
target_col = df['which_field'].map(field_map)
# 按行索引+列索引取值
df['output'] = df.values[df.index, df.columns.get_indexer(target_col)]

方案3：复杂自定义逻辑场景（灵活性最高）

如果判断逻辑不是简单的键值映射，需要多字段组合判断、自定义计算等，可使用apply实现：

def get_row_value(row):
    if row['which_field'] == 'a':
        return row['field_a']
    elif row['which_field'] == 'b':
        return row['field_b']
    # 可添加任意多判断逻辑
    # 示例：if row['field_a'] > 10 and row['which_field'] == 'c': return xxx
    else:
        return 0 # 自定义默认值

df['output'] = df.apply(get_row_value, axis=1)

运行以上任意方案都可以得到你预期的结果。

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