Remplir un tableau numpy 3D avec celles basées sur un tableau 2D avec des indices

Question

J'ai 2 tableaux de numpy output et index:

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

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

index représente l'indice jusqu'à ce que output devrait être peuplée avec ceux de la première dimension. La valeur peuplée de output devrait ressembler à:

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

Par exemple, depuis index[0, 1] == 1, nous avons mis en output[:1+1, 0, 1] = 1. En général, si index[i, j] == k, nous définissons output[:k+1, i, j] = 1.

Est-ce que quelqu'un sait comment y parvenir de manière vectorisée?

Answer 1

En utilisant NumPy broadcasting, nous pouvons créer un masque de ces endroits. Donc, convertir simplement ce masque dans un tableau int de 0s et 1s, comme si -

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

run échantillon -

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

Vous pouvez également assigner en output, utilisez le masque pour boolean-indexing et affecter 1s -

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

Answer 2

Vous pouvez attribuer 1 à la dernière position (le long de la première dimension) puis de nouveau remplir 0 avec 1 en utilisant np.maximum.accumulate:

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