如何使用idxmax结果筛选pandas DataFrame获取行最大值Series

问题背景 #

现有10万行、10列的随机数DataFrame，已通过以下代码得到逐行最大值max_series，以及每行最大值对应的列索引序列filter_：

import pandas as pd
import numpy as np

df = pd.DataFrame(np.random.random(size=(100000, 10)))
max_series = df.max(axis=1)
filter_ = df.idxmax(axis=1)

尝试通过df.loc[:, filter_]、df.filter(items=filter_, axis=1)提取值时触发内存错误，提示需要分配74.5GiB内存生成100000×100000的矩阵，和预期只输出10万长度的Series不符。

错误原因 #

上述两种写法的逻辑是按列筛选，传入长度为10万的filter_序列时，pandas会尝试匹配序列中所有值对应的列，最终要选10万列，自然会生成行数×列数=10万×10万的超大矩阵，触发内存溢出。

正确实现 #

最高效的写法是通过numpy的位置索引，逐行匹配对应列的取值，全程不会生成多余的中间矩阵，内存占用极低：

# 按行号、对应列号的位置索引直接取值，保留原df的索引
max_result = pd.Series(
    df.to_numpy()[np.arange(len(df)), filter_],
    index=df.index
)

执行后可以验证结果和直接调用max(axis=1)的结果完全一致：

print(max_result.equals(max_series))
# 输出: True

如果不想转numpy数组，也可以用pandas原生的逐行取值写法，性能略低于numpy方案但可读性更好：

max_result = df.agg(lambda row: row[filter_.loc[row.name]], axis=1)

不推荐使用apply/agg逐行循环的方案处理百万行以上规模的数据，numpy索引方案的速度是循环方案的上百倍。

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