基于其他列最小值索引为Pandas DataFrame新增列的高效实现

核心思路 #

针对50万行的大数据量，必须采用矢量化操作规避逐行处理的性能损耗，利用numpy的高效数组运算能力，直接完成批量索引匹配与取值。

步骤拆解：

• 先获取每行最小值对应的整数列索引（对应Seed1-Seed4的0-3顺序索引）
• 用该索引数组直接从RMI数组中批量取值，赋值为新列

代码实现 #

import numpy as np
import pandas as pd

# 初始化数据
RMI = [100, 200, 300, 400]
ar = np.array([[-2.0, 2.0, -5, 0], [2.7, 10, 5.4, 7], [5.3, 9, 1.5, -12]])
df = pd.DataFrame(ar, index=['1', '2', '3'], columns=['Seed1', 'Seed2', 'Seed3', 'Seed4'])

# 将RMI转为numpy数组，提升索引效率
rmi_arr = np.array(RMI)

# 获取每行最小值的整数列索引（axis=1表示按行计算）
min_col_indices = df.values.argmin(axis=1)

# 批量生成新列值
df['RMI_number'] = rmi_arr[min_col_indices]

print(df)

结果验证 #

输出结果完全符合预期：

Seed1  Seed2  Seed3  Seed4  RMI_number
1   -2.0    2.0   -5.0    0.0         300
2    2.7   10.0    5.4    7.0         100
3    5.3    9.0    1.5  -12.0         400

性能说明 #

• df.values.argmin(axis=1)：直接操作numpy底层数组，比idxmin（返回列名再映射索引）效率高数倍，尤其适配大数据量场景
• rmi_arr[min_col_indices]是numpy矢量化取值，时间复杂度为O(n)，远优于循环、apply等逐行处理方式

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