如何用Numpy加速二维数组按子区域聚合值的运算？

当你需要按子区域聚合值数组时，Numpy的向量化操作可以彻底替代低效的Python循环，下面是具体实现方案：

解决思路 #

核心是把二维数组扁平化后，利用Numpy内置的分组聚合函数（比如np.bincount）批量处理，避免逐区域迭代。这些函数都是底层C优化的，速度比Python循环快几个数量级。

代码示例 #

先假设我们有测试数据：

import numpy as np

# 模拟4x4的区域标记数组，区域ID为0-3
regions = np.array([
    [0,0,1,1],
    [0,2,2,1],
    [3,3,2,2],
    [3,3,3,3]
])
# 对应位置的数值数组
values = np.random.rand(4,4)

步骤1：扁平化数组 #

把二维的区域和值数组转成一维，方便后续操作：

regions_flat = regions.ravel()
values_flat = values.ravel()

步骤2：获取唯一区域ID #

用np.unique提取所有不重复的区域ID，同时拿到反向索引（用来快速把原数组元素映射到对应区域）：

unique_regions, idx_map = np.unique(regions_flat, return_inverse=True)

步骤3：执行聚合操作 #

求和聚合（最快方式）

np.bincount是处理这类分组求和的最优选择，直接按索引分组计算总和：

region_sums = np.bincount(idx_map, weights=values_flat)
# region_sums[i] 对应 unique_regions[i] 区域的总和

均值聚合

先统计每个区域的元素数量，再用总和除以数量：

region_counts = np.bincount(idx_map)
region_means = region_sums / region_counts

最大值/最小值聚合

如果需要求极值，因为np.bincount不支持，我们可以循环唯一区域ID（循环次数是唯一区域的数量，远少于原数组元素数）：

region_maxes = np.zeros(len(unique_regions), dtype=values.dtype)
for i, rid in enumerate(unique_regions):
    region_maxes[i] = values_flat[regions_flat == rid].max()

或者用列表推导式更简洁：

region_maxes = np.array([values_flat[regions_flat == rid].max() for rid in unique_regions])

为什么比循环快？ #

• Python循环是逐元素/逐区域的解释执行，而Numpy的操作是底层C实现的批量计算，避免了Python解释器的开销
• np.unique和np.bincount都是O(n)时间复杂度的向量化操作，处理百万级数据也能秒级完成

注意点 #

• 区域ID必须是整数类型，np.bincount只接受整数索引
• 如果区域ID不是连续整数，unique_regions会保留原始ID，聚合结果的顺序和它一一对应
• 对于超大数组（比如10000x10000），可以保持数组二维直接用np.add.at操作，避免扁平化的内存开销：

# 二维数组直接求和聚合的方式
unique_regions = np.unique(regions)
region_sums = np.zeros_like(unique_regions, dtype=values.dtype)
np.add.at(region_sums, idx_map, values)

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