Pandas按行选取最高值对应列并高效计算均值的实现方法

现有如下结构的Pandas DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame(np.array([[1, 2, 3, 3, 2, 1], [4, 3, 6, 6 ,3 ,4], [7, 2, 9, 9, 2, 7]]),
                   columns=['a', 'b', 'c', 'a_select','b_select','c_select'])

df

需要实现的计算逻辑：

• 逐行处理所有带_select后缀的列，筛选出每行值最大的2个_select列
• 匹配这2个列对应的不带_select后缀的同名列（例如a_select对应a列）
• 取匹配到的2个基础列在当前行的数值计算均值，注意计算均值的数据源是基础列，不是_select后缀列本身的数值
• 示例：第1行应对a、b列数值求均值，第2行应对a、c列数值求均值

当前采用逐行遍历的方式实现，写法简单但大数据集下运行速度很慢，需要用向量化方法实现等效逻辑提升效率。

不要使用iterrows遍历或者df.apply(axis=1)，这两种方式本质还是Python层面的行级循环，性能提升有限。直接基于numpy做向量化操作即可，性能比循环实现高1~2个数量级，百万级数据也能毫秒级出结果。

完整实现代码：

# 定义基础列和对应的筛选列，后续列数变动只需要修改base_cols即可
base_cols = ['a', 'b', 'c']
select_cols = [f"{col}_select" for col in base_cols]

select_arr = df[select_cols].to_numpy()
# 生成布尔掩码：标记每行select列中值最大的2个列的位置
top2_mask = select_arr >= np.sort(select_arr, axis=1)[:, [-2]]

# 用掩码过滤基础列数值，按行求均值
df["target_mean"] = (df[base_cols].to_numpy() * top2_mask).sum(axis=1) / 2

逻辑说明 #

1. 先拆分基础值列和筛选列，避免硬编码列名导致的匹配错误
2. 对筛选列的二维数组按行排序，取每行第二大的值作为阈值，生成布尔掩码：掩码为True的位置就是该行Top2的筛选列位置
3. 基础列数组和掩码相乘时，非选中位置的数值会被乘0失效，按行求和后除以2就是需要的均值

结果验证 #

运行后输出的target_mean列结果如下，完全符合预期：

0    1.5   # 第1行 (1+2)/2 = 1.5
1    5.0   # 第2行 (4+6)/2 = 5.0
2    8.0   # 第3行 (7+9)/2 = 8.0
Name: target_mean, dtype: float64

扩展说明 #

如果后续需要取Top N的列计算均值，只需要把代码中[-2]的索引改成[-N]，最后除数改成N即可，不需要调整核心逻辑。

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