如何使用NumPy加速Python循环优化大规模数组处理性能

云原生阿岚

2026-8-31

问题根因

你写的原生代码性能差的核心原因是采用了O(n²)的嵌套遍历逻辑：外层循环每处理一个元素，内层就要完整扫描15万元素的全量数组，总计算量达到225亿次，必然出现严重卡顿。

NumPy优化实现方案

优化核心思路是提前对相同area值的索引做分组聚合，把重复的全量扫描替换为O(1)的索引查询，再配合NumPy的向量化索引操作替代原生列表推导，整体时间复杂度可以降到O(n)级别。

完整实现代码

import numpy as np

# 1. 原生列表转NumPy数组
areas_arr = np.array(areas)
prices_arr = np.array(prices)
total_len = len(areas_arr)

# 2. 预构建 区域值->对应所有索引位置 的映射表，仅需一次全量遍历
val_index_map = {}
for idx, val in enumerate(areas_arr):
    if val not in val_index_map:
        val_index_map[val] = []
    val_index_map[val].append(idx)
# 映射表内的索引列表统一转NumPy数组，支持向量化筛选
for val in val_index_map:
    val_index_map[val] = np.array(val_index_map[val], dtype=np.int64)

# 3. 遍历处理每个元素，无需再全量扫描数组
for key in range(total_len):
    current_val = areas_arr[key]
    same_val_indices = val_index_map[current_val]
    # 向量化筛选排除当前索引，速度远快于原生列表判断
    comp_indices = same_val_indices[same_val_indices != key]
    compListPrices = prices_arr[comp_indices]
    
    # 在此处编写基于compListPrices的后续业务逻辑

额外优化提示

如果areas数组中存在大量唯一值，可以提前判断：若same_val_indices长度为1，说明没有其他同值元素，compListPrices为空数组，直接跳过后续计算即可。
如果后续对compListPrices的计算逻辑（比如求均值、中位数、求和等统计操作）可以用NumPy内置函数实现，尽量不要在Python层写循环处理，性能还能再提升1~2个数量级。