Numpy фантазии индексации в нескольких измерениях

Question

Допустим, у меня есть массив N размером n x m x k ​​и еще один массив B размером n x m с индексами от 1 до k. Я хочу получить доступ к каждому n x m срезу A, используя индекс, указанный в этом месте в B, давая мне массив размером п х м.

Редактировать: это, очевидно, не то, что я хочу! [[Я могу добиться этого, используя take, например:

A.take(B) ]] конец редактирования

Может ли это быть достигнуто с помощью необычной индексации? Я бы подумал, A[B] даст тот же результат, но это результаты в массиве размером n x m x m x k ​​(что я не очень понимаю).

Причина, по которой я не хочу использовать take, заключается в том, что я хочу иметь возможность назначить этой части что-то вроде

A[B] = 1

Единственное рабочее решение, которое у меня есть, это

A.reshape(-1, k)[np.arange(n * m), B.ravel()].reshape(n, m)

но наверняка должен быть более простой путь?

Answer 1

Пусть

import numpy as np
np.random.seed(0)

n,m,k = 2,3,5
A = np.arange(n*m*k,0,-1).reshape((n,m,k))
print(A)
# [[[30 29 28 27 26]
#   [25 24 23 22 21]
#   [20 19 18 17 16]]

#  [[15 14 13 12 11]
#   [10  9  8  7  6]
#   [ 5  4  3  2  1]]]

B = np.random.randint(k, size=(n,m))
print(B)
# [[4 0 3]
#  [3 3 1]]

Чтобы создать этот массив,

print(A.reshape(-1, k)[np.arange(n * m), B.ravel()])
# [26 25 17 12  7  4]

как массив nxm с использованием необычной индексации:

i,j = np.ogrid[0:n, 0:m]
print(A[i, j, B])
# [[26 25 17]
#  [12  7  4]]