如何使用NumPy where处理多维数组，替换不合格的2x2子数组？

问题背景 #

有一个形状为(N,2,2)的NumPy数组，包含N个2x2子数组；还有一个长度为N的布尔数组good，标记哪些子数组合格。需要将不合格的子数组替换为np.zeros((2,2))，使用np.where时出现广播错误：

import numpy as np

a = np.arange(20).reshape((5,2,2))
good = np.array([ x%4 != 3 for x in range(5) ])
np.where(good, a, np.zeros((2,2)))
# 报错：ValueError: operands could not be broadcast together with shapes (5,) (5,2,2) (2,2)

期望结果是第4个（索引3）子数组被替换为全0数组。

解决方案 #

方法1：调整布尔数组形状实现广播 #

报错原因是good的形状(5,)无法与a的(5,2,2)直接广播。只需给good增加两个维度，使其形状变为(5,1,1)，就能和a的后两个维度匹配：

# 方式1：用np.newaxis扩展维度
result = np.where(good[:, np.newaxis, np.newaxis], a, np.zeros((2,2)))

# 方式2：用reshape更简洁
result = np.where(good.reshape(-1,1,1), a, np.zeros((2,2)))

方法2：直接通过索引赋值（更高效） #

如果只需要将不合格子数组设为0，直接利用布尔索引赋值是最优方案，避免额外数组创建：

# 允许修改原数组的情况
a[~good] = 0

# 需要保留原数组的情况，先复制再修改
result = a.copy()
result[~good] = 0

方法3：利用广播创建全0数组 #

也可以先创建与a同形状的全0数组，再通过布尔索引保留合格子数组：

zeros = np.zeros_like(a)
result = np.where(good.reshape(-1,1,1), a, zeros)

说明 #

• 索引赋值方法（方法2）的时间和空间效率最高，因为它直接在数组上修改，无需生成中间数组。
• np.where方法更灵活，若后续需要将不合格子数组替换为其他非0值，只需修改第三个参数即可。

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