如何根据另一列存储的索引在Pandas中选取列？

下面几种方法都能快速实现你要的效果：

方法1：用numpy索引（高效推荐） #

直接借助numpy的索引能力，从A、B列的数值数组里按truth列的索引逐行取值：

import pandas as pd
import numpy as np

# 构建你的DataFrame
df = pd.DataFrame({
    'A': [123, 213, 312],
    'B': [132, 231, 321],
    'truth': [0, 1, 0]
})

# 生成output列
df['output'] = df[['A', 'B']].to_numpy()[np.arange(len(df)), df['truth']]

方法2：用df.lookup（旧版Pandas适用） #

lookup方法专门用来按行和列的索引匹配取值，步骤也很简单：

import pandas as pd

df = pd.DataFrame({
    'A': [123, 213, 312],
    'B': [132, 231, 321],
    'truth': [0, 1, 0]
})

# 先映射truth对应的列名，0对应'A'，1对应'B'
target_cols = df.columns[:2]
df['output'] = df.lookup(df.index, target_cols[df['truth']])

方法3：用apply（直观易读，适合小数据集） #

逐行遍历，根据truth的值选择对应列的内容：

import pandas as pd

df = pd.DataFrame({
    'A': [123, 213, 312],
    'B': [132, 231, 321],
    'truth': [0, 1, 0]
})

df['output'] = df.apply(lambda row: row[['A', 'B']][row['truth']], axis=1)

运行后都会得到你想要的结果：

A    B  truth  output
0  123  132      0     123
1  213  231      1     231
2  312  321      0     312

你之前用df.loc[]没效果是因为它更适合按标签或布尔条件批量选数据，没法处理这种逐行动态选列的需求，上面的方法都能完美解决问题。

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