python - 使用基于带有索引的 2D 数组的 1 填充 3D numpy 数组

Question

我有 2 个 numpy 数组输出和索引:

output = np.zeros((3,3,3))

>>>index   
array([[0,1,2],
       [1,0,0],
       [2,2,2]])

index 表示索引，直到 output 应该用第一维中的值填充。 output 的填充值应如下所示:

>>>output 
array([[[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]],
       [[0, 1, 1],
        [1, 0, 0],
        [1, 1, 1]],
       [[0, 0, 1],
        [0, 0, 0],
        [1, 1, 1]]] 

例如，由于 index[0, 1] == 1，我们设置 output[:1+1, 0, 1] = 1。一般来说，如果index[i, j] == k，我们设置output[:k+1, i, j] = 1。

有谁知道如何以矢量化方式实现这一目标？

Answer 1

使用NumPy broadcasting ，我们可以创建这些地方的掩模。因此，只需将该掩码转换为 0s 和 1s 的 int 数组，如下所示 -

(index >= np.arange(3)[:,None,None]).astype(int)

示例运行 -

In [471]: index
Out[471]: 
array([[0, 1, 2],
       [1, 0, 0],
       [2, 2, 2]])

In [472]: (index >= np.arange(3)[:,None,None]).astype(int)
Out[472]: 
array([[[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]],

       [[0, 1, 1],
        [1, 0, 0],
        [1, 1, 1]],

       [[0, 0, 1],
        [0, 0, 0],
        [1, 1, 1]]])

或者，要分配到输出，请使用 bool 索引的掩码并分配1s -

output[index >= np.arange(output.shape[0])[:,None,None]] = 1