Python四层嵌套循环的高效优化方法求助

你的核心需求是按纬度区间、经度区间、高度区间、小时区间分组计算温度均值，再映射到四维数组。四层Python循环效率极低，因为每次循环都要做全量数据的条件筛选，属于低效的逐组遍历操作。以下是两种基于向量化计算的高效实现方式：

方法一：Pandas 分箱 + GroupBy 聚合 #

利用Pandas的工具快速给每个数据点分配分组索引，再通过groupby一次性完成所有分组的均值计算，最后将结果填充到目标数组中。

步骤1：为每个字段生成分组索引 #

import pandas as pd
import numpy as np

# 纬度索引：对应lat_rng的区间（0到len(lat_rng)-2）
data['lat_bin'] = np.digitize(data['lat'], lat_rng) - 1
data['lat_bin'] = data['lat_bin'].clip(0, len(lat_rng)-2)  # 过滤超出范围的索引

# 经度索引：对应lon_rng的区间（0到len(lon_rng)-2）
data['lon_bin'] = np.digitize(data['lon'], lon_rng) - 1
data['lon_bin'] = data['lon_bin'].clip(0, len(lon_rng)-2)

# 高度索引：对应int_alt的中心区间（左闭右开）
alt_bins = np.concatenate([[int_alt[0]-gap_alt/2], int_alt + gap_alt/2])
data['alt_bin'] = np.digitize(data['alt'], alt_bins) - 1
data['alt_bin'] = data['alt_bin'].clip(0, len(int_alt)-1)

# 小时索引：LT的整数部分（0-23），处理LT=24的边界情况
data['lt_bin'] = data['LT'].astype(int)
data['lt_bin'] = data['lt_bin'].clip(0, 23)

步骤2：分组计算均值并填充到四维数组 #

# 按四个分组索引聚合，计算温度均值
grouped = data.groupby(['lon_bin', 'lat_bin', 'alt_bin', 'lt_bin'])['temp'].mean().reset_index()

# 初始化目标数组（和原代码的temp_all结构一致）
temp_all = np.full([len(lon_rng)-1, len(lat_rng)-1, len(int_alt), 24], np.nan)

# 将聚合结果填充到对应位置
for _, row in grouped.iterrows():
    lon_idx = int(row['lon_bin'])
    lat_idx = int(row['lat_bin'])
    alt_idx = int(row['alt_bin'])
    lt_idx = int(row['lt_bin'])
    temp_all[lon_idx, lat_idx, alt_idx, lt_idx] = row['temp']

方法二：Numpy 向量化分箱 + 加权均值（更适合超大数据量） #

如果数据量极大，可以用Numpy的底层操作实现聚合，避免Pandas的额外开销：

# 生成四个维度的索引（逻辑同方法一）
lat_bin = np.digitize(data['lat'], lat_rng) - 1
lat_bin = np.clip(lat_bin, 0, len(lat_rng)-2)
lon_bin = np.digitize(data['lon'], lon_rng) - 1
lon_bin = np.clip(lon_bin, 0, len(lon_rng)-2)
alt_bins = np.concatenate([[int_alt[0]-gap_alt/2], int_alt + gap_alt/2])
alt_bin = np.digitize(data['alt'], alt_bins) - 1
alt_bin = np.clip(alt_bin, 0, len(int_alt)-1)
lt_bin = np.clip(data['LT'].astype(int), 0, 23)

# 将四维索引转换为一维扁平索引（用于批量聚合）
flat_idx = lon_bin * (len(lat_rng)-1)*len(int_alt)*24 + \
           lat_bin * len(int_alt)*24 + \
           alt_bin *24 + \
           lt_bin

# 计算每个分组的温度总和与数据点数量
sum_temp = np.bincount(flat_idx, weights=data['temp'], minlength=temp_all.size)
count_temp = np.bincount(flat_idx, minlength=temp_all.size)

# 计算均值（避免除以0的情况）
mean_temp = np.where(count_temp > 0, sum_temp / count_temp, np.nan)

# 重塑为目标四维数组
temp_all = mean_temp.reshape(temp_all.shape)

性能优势说明 #

• 向量化计算：digitize、bincount都是C语言实现的底层操作，比Python循环快几个数量级
• 避免重复筛选：原循环每次都要对全量数据做条件判断，上述方法仅做一次分箱即可完成所有分组映射
• 批量处理：GroupBy或bincount一次性完成所有分组的均值计算，无需逐组遍历

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