Numpy：如何并行实现掩码应用后求均值（寻求更优方案）

我有一个一维Numpy数值数组：

import numpy as np
v = np.array([0, 1, 4, 0, 5])

此外，还有一个二维布尔掩码数组（实际场景中会有数百万个掩码）：

m = np.array([
    [True, True, False, False, False],
    [True, False, True, False, True],
    [True, True, True, True, True],
])

需要将掩码的每一行应用到数组v上，计算被掩码选中值的均值。预期效果如下：

results = []
for mask in m:
    results.append(np.mean(v[mask]))

print(results) # [0.5, 3.0, 2.0]

上述逻辑可通过循环实现，但我希望找到并行化的优雅实现方案。目前已找到一种方案：

mask = np.ones(m.shape)
mask[~m] = np.nan
np.nanmean(v * mask, axis=1) # [0.5, 3.0, 2.0]

是否存在其他方案（比如使用np.ma模块）？我正在寻找比当前两种方案更快的实现方式。

1. 使用np.ma模块实现 #

利用掩码数组直接对广播后的v做掩码处理，再计算均值：

v_broadcast = np.broadcast_to(v, m.shape)
masked_array = np.ma.masked_array(v_broadcast, mask=~m)
results = masked_array.mean(axis=1).data
# 输出：array([0.5, 3. , 2. ])

这种方法无需引入nan值，直接通过掩码忽略未选中元素，逻辑贴合需求，避免了浮点转换的潜在开销。

2. 求和+计数的高效实现（推荐百万级掩码场景） #

均值的本质是选中元素总和 / 选中元素数量，直接计算这两个值再做除法，是性能最优的方案：

sum_selected = np.dot(m, v)
count_selected = m.sum(axis=1)
results = sum_selected / count_selected
# 输出：array([0.5, 3. , 2. ])

这个方案的核心优势：

• 内存开销极小：无需将v广播到和m同形状，np.dot直接做矩阵-向量乘法，内存占用仅为广播方案的1/N（N为掩码行数）
• 计算效率极高：没有nan处理或掩码对象的额外逻辑，完全利用Numpy底层优化实现，速度远快于前两种方案，尤其适合百万级掩码的大规模场景。

性能对比参考 #

针对百万行的掩码数组：

• 循环方案：最慢，时间复杂度O(N*M)且无优化
• nanmean方案：中等，需广播和nan值处理
• np.ma方案：略优于nanmean，但仍有掩码对象的额外开销
• 求和计数方案：最快，矩阵乘法的底层优化让常数项极小，是大规模场景的首选

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