NumPy массив, индексирующий массив 4D

Question

У меня есть 4D массив 'a' размера (2,3,4,4), заполненный нулями.

import numpy as np 
a = np.zeros((2,3,4,4)) 

У меня также есть 3D-массива 'B' размера (2,3,4), который несет некоторые значения индекса (все от 0 до 3).

То, что я хочу сделать, это заменить элемент каждого последнего массива в «а» (4 измерения «а»), что соответствует индексу «Ъ», с 1.

я могу сделать это с 3 для петель, как показано ниже:

for i in a.shape[0]: 
    for j in a.shape[1]: 
     for z in a.shape[2]: 
      a[i,j,z][b[i,j,z]] = 1 

Но мне было интересно, если есть в любом случае я могу избежать зацикливания вообще. Нечто похожее на:

a[b] = 1 

Answer 1

Да, вы можете сделать это в векторизованной форме:

p,m,n,r = a.shape 
a.reshape(-1,r)[np.arange(p*m*n),b.ravel()] = 1 

Это следует обобщать более легко к более высокому ndarrays порядка.