如何为Pandas DataFrame每行按分类保留最优得分并优化性能?
高效处理大尺度Pandas DataFrame:按分类保留每行最高得分
针对你提到的250行(日期)、15000列(选手)、175个分类的场景,原转置-分组-idxmax-转置的方法会因大列数转置带来额外内存开销和性能损耗,以下是两种更高效的优化方案:
方案一:全向量化列分组Transform(推荐)
利用Pandas的列分组groupby(axis=1)结合transform,直接在原DataFrame结构上计算每行各分类的最大值,再通过掩码保留最大值位置,其余置0/NaN:
import pandas as pd import numpy as np # 假设df为原始得分DataFrame,cat_map是列名(选手)到分类的映射Series # 1. 按分类分组列,计算每行的分类最大值(自动忽略NaN) group_max = df.groupby(cat_map, axis=1).transform('max') # 2. 生成掩码:仅保留原DataFrame中等于对应分类最大值的位置 mask = df == group_max # 3. 保留最大值,其余置0(如需置NaN,替换0为np.nan即可) result = df.where(mask, 0)
优势:
- 完全向量化操作,无转置、无循环,Pandas内部优化程度高
- 避免转置带来的内存复制(15000列转置后变为15000行,内存开销翻倍)
- 代码简洁,可读性强,处理250行的小行数场景速度极快
方案二:分类循环+向量化掩码(内存友好)
如果服务器内存紧张,可按分类循环处理,每次仅操作单个分类的列子集,进一步降低内存占用:
import pandas as pd import numpy as np # 初始化结果DataFrame为全0(或全NaN) result = pd.DataFrame(0, index=df.index, columns=df.columns) # 遍历每个分类 for cat in cat_map.unique(): # 提取当前分类对应的所有选手列 cat_cols = cat_map[cat_map == cat].index cat_data = df[cat_cols] # 计算当前分类每行的最大值 cat_row_max = cat_data.max(axis=1) # 生成掩码:标记当前分类中等于该行最大值的位置 mask = cat_data.eq(cat_row_max, axis=0) # 将最大值写入结果DataFrame对应位置 result.loc[:, cat_cols] = result.loc[:, cat_cols].where(~mask, cat_data)
优势:
- 每次仅加载单个分类的列(15000/175≈85列),内存占用远低于方案一
- 内部仍采用向量化操作,循环仅针对175个分类,开销可忽略
性能对比
原转置方案的瓶颈在于大列数转置的内存复制和大行数分组的计算开销,而上述两种方案均在原DataFrame的行维度上操作,利用Pandas的向量化优化,在你的场景下速度可提升3-10倍(具体取决于内存带宽)。
内容的提问来源于stack exchange,提问作者lubenthrust
相关产品推荐
相关产品推荐

