根据具有索引的2D数组填充3D numpy数组

Question

我有2个numpy数组output和index ：

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的整数数组，就像这样-

(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]]])

或者，要分配给output ，请使用掩码进行boolean-indexing并分配1s

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

Answer 2

您可以将1分配给最后一个位置（沿第一个维度），然后使用np.maximum.accumulate用1 np.maximum.accumulate 0：

output[index, np.arange(output.shape[1])[:,None], np.arange(output.shape[2])] = 1  
np.maximum.accumulate(output[::-1], axis=0)[::-1]

#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.]]])