如何高效基于多标签掩码对NumPy数组按标签分组求和

架构老林

2026-8-30

高效按标签分组求和实现方法

你当前的循环实现时间复杂度为O(K*N)（K为标签总数，N为数组元素总数），每次迭代都要全量扫描标签数组生成掩码、再索引取值，数据规模变大时性能会线性下降。可以使用单趟遍历的原生API实现，时间复杂度降到O(N)，性能提升可达两个数量级以上。

方案1：无额外依赖，numpy原生`np.bincount`（性能最优）

这是纯numpy场景下的最快实现，核心思路是用标签作为分桶索引，直接对对应位置的x值做累加，全程仅需遍历一次数组：

import numpy as np

x = np.array([[0, 1, 2],
              [3, 4, 5],
              [6, 7, 8]])
labels = np.array([[0, 0, 2],
                   [1, 1, 2],
                   [1, 1, 2]])

# 核心计算逻辑，一行完成
out = np.bincount(labels.ravel(), weights=x.ravel())

print(out)
# 输出: [ 1. 20. 15.]

说明

ravel()方法会把二维数组展平为一维，不会产生数据拷贝（除非数组内存不连续），开销极低
注意：该方法要求标签是从0开始的连续非负整数，如果你的标签存在跳号、负数，可以先用np.unique(labels, return_inverse=True)把标签映射为连续编号后再计算
如果需要整数类型结果，直接在结果后调用.astype(np.int64)即可

方案2：支持非连续标签/无效值，scipy.ndimage 实现

如果你的标签不连续，或者存在需要跳过的无效标记（比如标记为-1的背景区域），可以用scipy的ndimage模块提供的聚合函数，不需要手动做标签映射：

from scipy import ndimage

# index参数传入你需要统计的标签列表即可
out = ndimage.sum(x, labels=labels, index=np.arange(labels.max() + 1))

print(out)
# 输出: [ 1. 20. 15.]

该方法内部同样是单趟遍历实现，性能和np.bincount接近，额外支持自定义统计标签范围、忽略指定无效标签值，适合复杂掩码场景。