如何用广播布尔掩码索引数组，避免创建巨型中间数组？

问题分析 #

直接通过广播a1[..., None]和a2[None, ...]匹配掩码维度时，会生成M1×…×MP×N1×…×NQ×3的巨型中间数组，高维度场景下内存占用会急剧飙升。我们可以通过提取掩码对应位置的索引，直接索引原数组来规避这个问题。

解决方案代码 #

import numpy as np

np.random.seed(0)

a1 = np.random.rand(4, 5, 3)
a2 = np.random.rand(6, 3)
m = np.random.rand(4, 5, 6) >= 0.7

# 获取掩码为True的位置索引
indices = np.where(m)

# 拆分索引：前P个维度对应a1的索引（P = a1.ndim - 1，排除最后一维的3）
a1_indices = indices[:a1.ndim - 1]
b1 = a1[a1_indices]

# 拆分索引：后Q个维度对应a2的索引（Q = a2.ndim - 1，排除最后一维的3）
a2_indices = indices[a1.ndim - 1:]
b2 = a2[a2_indices]

原理说明 #

1. np.where(m)返回掩码中所有True位置的多维索引元组，每个元素是对应维度的索引数组，内存占用远小于广播生成的巨型数组。
2. 拆分索引时，a1的有效维度是除最后一维（3）外的前P个维度，对应索引元组的前P个元素；a2的有效维度是除最后一维（3）外的后Q个维度，对应索引元组的剩余元素。
3. 直接用拆分后的索引访问原数组，无需生成中间广播数组，大幅降低内存消耗。

验证等价性 #

如果需要验证优化后的结果与原始高内存方法的一致性，可以运行以下代码：

# 高内存的原始实现（仅用于验证）
b1_original = a1[..., None].repeat(a2.shape[0], axis=-2)[m, :]
b2_original = a2[None, ...].repeat(a1.shape[0], axis=0).repeat(a1.shape[1], axis=1)[m, :]

print(np.allclose(b1, b1_original))  # 输出 True
print(np.allclose(b2, b2_original))  # 输出 True

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