pandas中使用分类器输出筛选DataFrame每行对应列元素的实现方案

实现方案 #

方法1：使用NumPy高级索引（性能最优，适合大数据量场景） #

直接利用numpy数组的花式索引特性，一次性匹配所有行对应的列位置，时间复杂度为O(n)，性能远高于逐行遍历：

import pandas as pd
import numpy as np

data = {'zero':[1,2,3], 'one':[4,5,6], 'two':[7,8,9]}
df = pd.DataFrame(data)
clf = [1,0,2]

# 核心代码
output = df.values[np.arange(len(df)), clf].tolist()
print(output) # 输出 [4, 2, 9]

原理说明：np.arange(len(df))生成和行数等长的行索引序列，和clf的列索引序列一一对应，直接从DataFrame对应的numpy数组中提取对应位置的元素。

方法2：使用apply逐行提取（写法直观，适合小数据量场景） #

如果数据量不大，追求写法的可读性，可以用pandas的行遍历方法实现：

output = df.apply(lambda row: row.iloc[clf[row.name]], axis=1).tolist()
print(output) # 输出 [4, 2, 9]

原理说明：row.name获取当前行的索引位置，匹配clf中对应位置的列索引，提取该行对应列的元素。

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