如何根据行内存储的列名字段匹配提取对应列值生成DataFrame新列

实现方案 #

下面给出两种不同场景下的实现方式，均可以满足需求：

1. 中小数据量场景（可读性优先） #

用apply逐行处理，代码逻辑直观易懂：

import pandas as pd

# 构造示例数据
df = pd.DataFrame({
    "col1": [1, 0],
    "col2": [1, 0],
    "col3": [0, 1],
    "col4": [1, 1],
    "colfromvaluestobepicked": ["col1", "col2"]
})

# 按规则生成新列
df["new col"] = df.apply(lambda row: row[row["colfromvaluestobepicked"]], axis=1)

2. 大数据量场景（性能优先） #

apply本质是行循环，数据量超过10万行时性能会明显下降，推荐使用向量化方案，比apply快10~100倍：

import pandas as pd
import numpy as np

# 构造示例数据同上，此处省略
idx, cols = pd.factorize(df["colfromvaluestobepicked"])
df["new col"] = df.reindex(cols, axis=1).to_numpy()[np.arange(len(df)), idx]

如果使用pandas 2.2版本以下，也可以用更简洁的lookup方法实现：

df["new col"] = df.lookup(df.index, df["colfromvaluestobepicked"])

两种方案运行后得到的结果和你示例中给出的new col完全一致。

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