Python优化DataFrame指定列行内最大值更新效率问题求助

问题描述 #

原始DataFrame如下：

Day US  INDIA   JAPAN   GERMANY AUSTRALIA Threshold
11  40  30      20      100     110         5
21  60  70      80      55      57          8
32  12  43      57      87      98          9
41  99  23      45      65      78          12

需求：在每行的US、INDIA、GERMANY三列中找出最大值，将该值与对应行的Threshold相加后，更新回原DataFrame的对应位置，期望结果：

Day US  INDIA   JAPAN   GERMANY AUSTRALIA Threshold
11  40  30      20      105     110       5
21  60  78      80      55      57        8
32  12  43      57      96      98        9
41  111 23      45      65      78        12

当前使用for循环实现，但大数据量下执行速度过慢，循环代码：

df_max = df_final[['US','INDIA','GERMANY']].idxmax(axis=1)
for ind in df_final.index:
    column = df_max[ind]
    df_final[column][ind] = df_final[column][ind] + df_final['Threshold'][ind]

高效解决方案 #

利用Pandas的矢量化操作替代循环，底层优化的矢量化操作能大幅提升处理速度，代码如下：

import pandas as pd

# 构造示例数据（实际使用时替换为你的df_final）
data = {
    'Day': [11, 21, 32, 41],
    'US': [40, 60, 12, 99],
    'INDIA': [30, 70, 43, 23],
    'JAPAN': [20, 80, 57, 45],
    'GERMANY': [100, 55, 87, 65],
    'AUSTRALIA': [110, 57, 98, 78],
    'Threshold': [5, 8, 9, 12]
}
df_final = pd.DataFrame(data)

# 定义需要处理的目标列
target_cols = ['US', 'INDIA', 'GERMANY']

# 生成每行最大值的位置掩码：True表示该位置是对应行的最大值
max_mask = df_final[target_cols] == df_final[target_cols].max(axis=1).values[:, None]

# 仅对最大值位置的数值加上Threshold，其余位置保持不变
df_final[target_cols] = df_final[target_cols].where(
    ~max_mask,  # 掩码取反，保留非最大值的原始值
    df_final[target_cols] + df_final['Threshold'].values[:, None]  # 最大值位置加上Threshold
)

# 输出结果
print(df_final)

方案说明 #

1. 生成最大值掩码：通过df_final[target_cols].max(axis=1)获取每行的最大值，再与原列比较生成布尔掩码，标记出每行最大值的位置。
2. 批量更新数值：使用where函数，仅对掩码为True的位置（即最大值位置）执行数值加法操作，全程无需循环，完全利用Pandas的矢量化计算优势。

这种方法的时间复杂度远低于循环方案，尤其在处理十万级以上数据时，性能提升会非常明显。

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