如何用纯NumPy函数实现按布尔掩码分组计算数组均值

解决方案 #

直接利用NumPy的矩阵乘法和广播机制，一次性完成所有分组的均值计算，彻底消除Python循环的开销：

import numpy as np

# 计算每个分组的有效元素数量
group_counts = indices.sum(axis=1, keepdims=True)
# 处理空分组（避免除以0的情况）
group_counts[group_counts == 0] = 1
# 矩阵乘法求分组总和，再除以元素数得到均值
means = (indices @ data) / group_counts

性能提升原理 #

• 原列表推导式会循环30次，每次执行data[i]都会生成临时子数组，带来大量内存和计算开销。
• indices @ data是底层C实现的矩阵运算，直接对30个分组的行求和，输出形状为(30, 3650)的总和矩阵。
• group_counts保持(30,1)的维度，通过广播机制自动和总和矩阵的每列做除法，高效得到均值。

空分组处理优化 #

如果存在全False的掩码（空分组），可以将对应均值设为NaN更合理：

group_counts = indices.sum(axis=1, keepdims=True)
means = np.divide(indices @ data, group_counts, out=np.full_like(indices @ data, np.nan), where=group_counts != 0)

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