如何用Numpy基于2D深度数组生成3D二进制掩码？

核心思路是利用Numpy的广播机制，通过维度扩展让2D深度数组与第三维度的索引数组做逐元素比较，直接生成目标掩码，全程无需循环。

实现步骤 #

• 生成第三维度的索引数组，对应每个位置的索引0到8（与输出数组的第三维度长度一致）
• 将2D深度数组扩展为3D（新增一个维度），利用广播与索引数组比较，得到布尔掩码
• 将布尔数组转换为数值类型（float或int），得到最终的3D掩码数组

代码示例 #

import numpy as np

# 定义输入的2D深度数组
tst = np.arange(1, 10).reshape(3,3)
# 生成第三维度的索引
depth_indices = np.arange(9)
# 利用广播生成掩码数组
blk = (depth_indices < tst[..., np.newaxis]).astype(float)

结果验证 #

print(blk[0,0])
# 输出：array([1., 0., 0., 0., 0., 0., 0., 0., 0.])

print(blk[1,2])
# 输出：array([1., 1., 1., 1., 1., 1., 0., 0., 0.])

print(blk[2,2])
# 输出：array([1., 1., 1., 1., 1., 1., 1., 1., 1.])

补充说明 #

tst[..., np.newaxis]的作用是将原(3,3)的2D数组扩展为(3,3,1)的3D数组，这样Numpy会自动将depth_indices（形状(9,)）广播为(3,3,9)，与扩展后的tst逐元素比较，得到的布尔数组直接对应需要填充1或0的位置，最后转换为数值类型即可。

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