如何高效获取DataFrame每行中n个最大值对应的列名？

现有如下DataFrame数据：

import pandas as pd
import numpy as np

np.random.seed(42)
data = np.random.rand(5,5)
data = pd.DataFrame(data, columns = list('abcde'))

对应的DataFrame内容为：

a         b         c         d         e
0  0.374540  0.950714  0.731994  0.598658  0.156019
1  0.155995  0.058084  0.866176  0.601115  0.708073
2  0.020584  0.969910  0.832443  0.212339  0.181825
3  0.183405  0.304242  0.524756  0.431945  0.291229
4  0.611853  0.139494  0.292145  0.366362  0.456070

需要提取每行中数值最大的n个元素对应的列名（以n=2为例），期望输出如下：

0  b  c
1  c  e
2  b  c
3  c  d
4  a  e

不想通过循环遍历实现（效率低），求更Pythonic的高效实现方式。

方法一：Pandas原生简洁实现 #

利用DataFrame.apply()结合Series.nlargest()，直接提取每行前n大值的列名：

n = 2
# 提取每行前n大值的列名并转为DataFrame
result = data.apply(lambda row: row.nlargest(n).index.tolist(), axis=1).apply(pd.Series)
print(result)

输出结果与期望完全一致，该方法代码简洁直观，适合中小型数据集。

方法二：Numpy向量化操作（高效处理大数据） #

当数据集规模较大时，使用Numpy的向量化操作可以避免Python层面的循环，性能优势更明显：

n = 2
# 获取每行前n大值的索引（argpartition仅保证前n个为最大，顺序未排序）
top_indices = np.argpartition(-data.values, n, axis=1)[:, :n]
# 对前n个索引按对应数值从大到小排序
sorted_indices = top_indices[np.arange(len(top_indices))[:, None], 
                             np.argsort(-data.values[np.arange(len(top_indices))[:, None], top_indices])]
# 将索引映射为列名，转为DataFrame
result = pd.DataFrame(data.columns[sorted_indices], index=data.index)
print(result)

该方法依托Numpy的底层优化，处理百万级以上数据时效率显著高于Pandas的apply()方法。

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