pandas气象预测DataFrame比较列值与Location分组降雨量均值替换问题

原代码问题 #

• 分组聚合得到的df_rainfall索引为Location值，本身不存在Location列，调用df_rainfall['Location']会直接触发KeyError；即使重置索引后，df_rainfall仅44行，和原数据集11万+行的Location列长度完全不匹配，equals判断永远无法成立
• 循环内的result是临时变量，赋值操作不会修改原DataFrame的实际存储值，即使逻辑正确也无法生效
• 逐行/逐列循环属于pandas不推荐用法，处理10万级以上数据性能极低

最优实现代码 #

使用pandas的transform方法做向量化操作，无需循环，瞬间完成10万+数据计算：

import numpy as np

# 生成分组均值序列，长度、索引和原数据集完全对齐，每个位置对应当前行所属Location的平均降雨量
location_rain_mean = weather_train_data_total.groupby('Location')['Rainfall'].transform('mean')
# 逐元素取较小值：降雨量大于均值时替换为均值，否则保留原值
weather_train_data_total['Rainfall'] = np.minimum(weather_train_data_total['Rainfall'], location_rain_mean)

如果不想依赖numpy，可以用pandas原生where方法实现相同逻辑：

location_rain_mean = weather_train_data_total.groupby('Location')['Rainfall'].transform('mean')
weather_train_data_total['Rainfall'] = weather_train_data_total['Rainfall'].where(
    cond = weather_train_data_total['Rainfall'] <= location_rain_mean,
    other = location_rain_mean
)

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