Pandas如何免嵌套循环按指定行列索引提取值生成同维度列表

核心思路是利用numpy的花式索引一次性批量提取所有目标值，避免Python层嵌套循环的取值开销，数据量越大性能优势越明显。

具体实现代码 #

import numpy as np
import pandas as pd

# 给定输入数据
names = [['a','b'],['c','c'],['b','c']]
df = pd.DataFrame({
    'a': [1,1,4],
    'b': [2,3,6],
    'c': [6,2,4]
})

# 批量取数逻辑
row_ids = np.arange(len(names)).reshape(-1, 1)  # 每个子列表对应df的行下标，做维度扩展用于广播
col_ids = np.array([df.columns.get_indexer(cols) for cols in names])  # 批量将列名转为df对应的整数列位置
lst = df.to_numpy()[row_ids, col_ids].tolist()

运行后得到的结果和原嵌套循环逻辑完全一致：

print(lst)
# 输出: [[1, 2], [2, 2], [6, 4]]

逻辑说明 #

• 扩展维度后的行索引会自动和列索引做广播匹配，严格对应names第i个子列表取df第i行的规则
• 列索引转换一次性完成所有列名到位置的映射，不需要逐行逐列在df里取值
• 所有取值操作在numpy底层完成，最终调用tolist()直接输出和names维度完全一致的嵌套列表

如果要完全规避Python层显式循环，也可以用itertools拍平列名后做维度转换，性能差异很小：

from itertools import chain
row_ids = np.repeat(np.arange(len(names)), [len(cols) for cols in names])
col_ids = df.columns.get_indexer(chain.from_iterable(names))
lst = df.to_numpy()[row_ids, col_ids].reshape(len(names), -1).tolist()

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