Использовать массив как индексную маску для многомерного массива

Question

У меня есть следующие массивы:

a = np.arange(12).reshape((2, 2, 3))

и

b = np.zeros((2, 2))

Теперь я хочу использовать b для доступа к a, st на каждом для индекса i, j мы берем z-ый элемент a, если b [i, j] = z.Для приведенного выше примера ответом должно быть [[0, 3], [6, 9]].Я чувствую, что это очень связано с np.choose, но все же как-то не совсем справляется.Вы можете мне помочь?

Answer 1

Можно предложить два подхода.

С массивами диапазона explicit для advanced-indexing -

m,n = b.shape
out = a[np.arange(m)[:,None],np.arange(n),b.astype(int)]

С np.take_along_axis -

np.take_along_axis(a,b.astype(int)[...,None],axis=2)[...,0]

Пробный прогон -

In [44]: a
Out[44]: 
array([[[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 6,  7,  8],
        [ 9, 10, 11]]])

In [45]: b
Out[45]: 
array([[0., 0.],
       [0., 0.]])

In [46]: m,n = b.shape

In [47]: a[np.arange(m)[:,None],np.arange(n),b.astype(int)]
Out[47]: 
array([[0, 3],
       [6, 9]])

In [48]: np.take_along_axis(a,b.astype(int)[...,None],axis=2)[...,0]
Out[48]: 
array([[0, 3],
       [6, 9]])