在 numpy 中使用 2d 深度值数组创建 3d 二进制掩码

Question

我想基于“深度”的 2d 数组创建一个 3d 掩码数组，该数组指示第三维填充的程度。我本来希望 numpy 中有一个向量化的解决方案，而不是诉诸 for 循环。 例如：

tst = np.arange(1, 10).reshape(3,3) # depths referred to above

blk = np.zeros(shape=(3,3,9)) # the output 3d array to be filled with ones or zeros

以下示例位置的预期结果是：

blk[0,0]
array([1., 0., 0., 0., 0., 0., 0., 0., 0.])
blk[1,2]
array([1., 1., 1., 1., 1., 1., 0., 0., 0.])
blk[2,2]
array([1., 1., 1., 1., 1., 1., 1., 1., 1.])

有人对此有任何想法/提示吗？

Answer 1

可以使用广播，相比

numpy.arange

：

blk = (tst[..., None] > np.arange(tst.max())[None, None]).astype(int)

输出：

array([[[1, 0, 0, 0, 0, 0, 0, 0, 0],     # blk[0, 0]
        [1, 1, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0, 0]],

       [[1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 0, 0, 0]],    # blk[1, 2]

       [[1, 1, 1, 1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1]]])   # blk[2, 2]