17370845950

本文介绍一种高效、向量化的方法，利用 `pd.factorize` 和 numpy 高级索引，根据辅助 dataframe 中的列名字符串，从主 dataframe 中按行提取对应列的值，适用于大规模数据场景。

在实际数据分析中，常遇到这样的需求：一个 DataFrame（如 df）存储多列数值数据，另一个 DataFrame（如 df1）的某列（如 'idx'）存储列名字符串（如 "a" 或 "b"），要求为每一行动态选取 df 中同索引、且列名为 df1['idx'] 对应值的那一列元素。例如：

import pandas as pd
import numpy as np

df = pd.DataFrame({'a': [94, 170, 5],
                   'b': [31, 115, 8]}, index=[11, 12, 13])

df1 = pd.DataFrame({'idx': ["a", "b", "a"]}, index=[11, 12, 13])

目标是得到结果 [94, 115, 5] —— 即第 11 行取 'a' 列（94），第 12 行取 'b' 列（115），第 13 行取 'a' 列（5）。

✅ 推荐方案：reindex + factorize + NumPy 高级索引（高性能）

该方法完全向量化，避免 apply 或 Python 循环，适合百万级数据：

idx, cols = pd.factorize(df1['idx'])  # 将列名映射为整数编码（如 'a'→0, 'b'→1）

# 重索引 df：确保行索引对齐 df1.index，列只保留 cols 中出现的列（安全且高效）
aligned_df = df.reindex(index=df1.index, columns=cols)

# 转为 NumPy 数组，用 np.arange(len(df)) 构造行索引，idx 构造列索引，实现逐行“列名定位”
result = aligned_df.to_numpy()[np.arange(len(df1)), idx]

print(result)  # [ 94 115   5]

若需将结果作为新列加入 df1：

df1['out'] = result
# 或一步写成：
df1['out'] = (df.reindex(index=df1.index, columns=cols)
                .to_numpy()[np.arange(len(df1)), idx])

输出：

   idx  out
11   a   94
12   b  115
13   a    5

⚠️ 注意事项与最佳实践

• 索引必须对齐：df 和 df1 的索引需一致（或至少 df1.index 是 df.index 的子集），否则 reindex 会引入 NaN；
• 列名存在性校验：pd.factorize 不检查 df1['idx'] 中的列名是否真实存在于 df.columns。建议提前校验：

  invalid_cols = set(df1['idx']) - set(df.columns)
  if invalid_cols:
      raise ValueError(f"Column names not found in df: {invalid_cols}")

• 内存优化：reindex(..., columns=cols) 可显著减少内存占用，尤其当 df 列数远多于 df1['idx'] 实际引用的列时；
• 替代方案对比：
  • ❌ df.lookup() 已在 Pandas 2.0+ 中弃用，不应使用；
  • ❌ df.apply(lambda x: df.loc[x.name, df1.loc[x.name, 'idx']], axis=1) 效率极低，为 O(n) Python 循环；
  • ✅ 本方案时间复杂度接近 O(n)，底层调用 NumPy C 实现，实测在百万行上比 apply 快 100+ 倍。

✅ 总结

当需基于列名动态取值时，优先采用 factorize + reindex + NumPy 高级索引 组合。它兼具正确性、可读性与极致性能，是处理大规模“列名驱动索引”任务的标准范式。务必确保索引对齐与列名有效性，即可安全应用于生产环境。