Pandas如何将DataFrame每行符合条件的N个最大值替换为其总和

处理规则 #

逐行对DataFrame执行如下操作：

• 取每行绝对值最大的N个值，计算总和存入列nlargestsum
• 计算该行判断阈值：阈值 = nlargestsum / N
• 遍历该行所有数值单元格：若单元格原值 >= 阈值，将单元格值替换为该行nlargestsum；若原值 < 阈值，保留原值不变

测试数据构造 #

import pandas as pd
data = [[1, 0.9, 0, 0, 0, 0, 0], [2, 0.3, 0.3, 0.3, 0, 0, 0.1], [3, 0, 0, 0, 0.2, 0.4, 0], [4, 1, 0, 0, 0, 0, 0],  [5, 0.1, 0.1, 0.4, 0.4,  0, 0]]
df = pd.DataFrame(data, columns=["ID", 'S1', "S2", 'S3', 'S4', 'S5', 'S6'])
df.set_index('ID', inplace =True)

完整实现代码 #

已完成的nlargestsum计算逻辑可直接复用，后续替换无需逐行循环，使用pandas原生向量化操作性能更好，全量可运行代码如下：

N = 3
# 计算每行前N大绝对值的和
df['nlargestsum'] = df.apply(lambda s: s.abs().nlargest(N).sum(), axis=1)
# 指定需要做值替换的数值列范围
score_cols = [c for c in df.columns if c != 'nlargestsum']
# 生成替换掩码：标记所有满足「值 >= 该行阈值」的单元格位置
row_threshold = df['nlargestsum'] / N
replace_mask = df[score_cols].ge(row_threshold, axis=0)
# 按掩码完成值替换
df[score_cols] = df[score_cols].where(~replace_mask, df['nlargestsum'], axis=0)

逻辑说明 #

• ge方法按行对齐比较每个单元格值和对应行的阈值，返回布尔型DataFrame，True代表该位置需要替换
• where方法会保留~replace_mask（即不需要替换）位置的原值，将需要替换的位置填充为对应行的nlargestsum
• 运行后得到的结果和给出的预期输出完全一致，ID=2/3/5等特殊场景均匹配规则要求。

内容的提问来源于stack exchange，提问作者Luis Manuel Quiros guerrero