如何无循环实现Numpy三维数组与二维0-1掩码的类按位与操作

现有形状为(172,40,20)的NumPy三维数组，以及形状为(172,40)的0-1掩码数组，需要实现类似按位与的操作：保留掩码值为1对应位置的原数组值，掩码值为0的位置设为0，且不能用循环实现。示例如下：

import numpy as np

a3D = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
print(a3D.shape)  # 输出：(2, 2, 2)
mask = np.array([[1, 0], [0, 1]])
print(mask.shape)  # 输出：(2, 2)

# 期望得到的结果
# result = np.array([[[1, 2], [0, 0]], [[0, 0], [7, 8]]])

利用NumPy的广播机制即可实现，不需要循环。核心思路是给掩码数组增加一个维度，让它和三维数组的维度对齐，然后直接做乘法运算（因为掩码是0-1值，乘法等价于保留1对应位置、置0对应位置）。

具体实现代码：

import numpy as np

def bitwise_and_3d_mask(a3D, mask):
    # 给掩码增加最后一个维度，形状从(172,40)变为(172,40,1)，触发广播
    return a3D * mask[:, :, np.newaxis]

# 测试示例
a3D = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
mask = np.array([[1, 0], [0, 1]])
result = bitwise_and_3d_mask(a3D, mask)
print(result)
# 输出：
# [[[1 2]
#   [0 0]]
# 
#  [[0 0]
#   [7 8]]]

原理说明 #

• NumPy广播会自动将形状为(172,40,1)的掩码数组扩展为(172,40,20)，和原三维数组维度匹配；
• 0-1掩码与原数组相乘时，掩码为1的位置保留原数组值，掩码为0的位置直接置0，完全符合需求；
• 这种方式是NumPy原生的向量化操作，效率远高于循环实现。

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