如何加速仅用Pandas内置函数的apply操作？

问题场景 #

现有如下DataFrame df：

当前需求是：按行提取前两大值对应的列名，再取这些列名的最大值，使用的代码为：

df.apply(lambda r: r.nlargest(2).index.max(), axis=1)

优化方案：使用Numpy向量化操作摆脱Python循环 #

apply本质是Python级别的逐行循环，数据量较大时性能瓶颈明显。可以利用Numpy的argpartition实现O(n)时间复杂度的向量化操作，大幅提升速度：

步骤1：分离数值列与日期列 #

import numpy as np
import pandas as pd

# 分离数值列和日期列（假设交易日期为非数值列）
numeric_cols = df.columns.drop('交易日期')
values = df[numeric_cols].values
col_names = numeric_cols.to_numpy()

步骤2：快速定位每行前两大值的位置 #

argpartition能在不完整排序的情况下，快速将前k个最大元素移到数组前端，效率远高于全排序：

# 获取每行前两大值的索引位置
top2_indices = np.argpartition(-values, 1, axis=1)[:, :2]

步骤3：提取对应列名并取最大值 #

# 提取前两大值对应的列名，再按行取最大值
result = col_names[top2_indices].max(axis=1)

# 若需要与原DataFrame的交易日期对应，转为Series
result_series = pd.Series(result, index=df['交易日期'], name='max_top2_col')

性能说明 #

• argpartition是基于分区的算法，时间复杂度为O(n*m)（n为行数，m为列数），但实际执行效率远高于apply的逐行循环，尤其是当数据量达到十万级以上时，性能差距会非常显著。
• 若坚持使用Pandas原生方法，可尝试用rank标记前两大值的列，但本质仍为apply循环，性能提升有限：

rank_df = df[numeric_cols].rank(axis=1, ascending=False, method='min') <= 2
result = rank_df.apply(lambda x: x.index[x].max(), axis=1)

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