Извлечение значений из последнего измерения трехмерного массива

Question

Я пытаюсь извлечь из трехмерной матрицы значений двухмерную матрицу, в которой последнее измерение имеет значения из последнего измерения трехмерной матрицы.Например, если P измерений [2,2,3] =

[ [[5, 1, 5], [9, 9, 4]], [[0, 9, 8], [8, 6, 8]] ]

, что такое индексная матрица, чтобы получить выходную матрицу [[1, 9],[0, 8]], где 1 - 2-й элементпервого столбца первого ряда, 9 - это 1-й элемент второго столбца первого ряда, 0 - 1-й элемент первого столбца второго ряда, а 8 - 3-й элемент второго столбца второго ряда?Идея состоит в том, что для каждого столбца у меня есть k разных баллов.Я хочу получить для каждого столбца разные оценки, для которых я знаю индекс.

Я немного запутался в продвинутой индексации в Numpy, и сам по себе это не понимаю.Спасибо!

Answer 1

Я предполагаю, что есть индексный массив для индексации по последней оси.Давайте назовем это idx.Для данного образца с данным текстом в вопросе это будет -

idx = np.array([[1,0],[0,2]])

В частности, это было извлечено из цитируемого текста:

1 - это 2-й элемент первого столбца первого ряда, 9 - 1-й элемент второго столбца первого ряда, 0 - 1-й элемент первого столбца второго ряда и8 - это 3-й элемент второго столбца второй строки

. Для решения этой задачи мы будем использовать открытую сетку с np.ogrid для индексации первых двух осей входного массива -

m,n = idx.shape
I,J = np.ogrid[:m,:n]
out = A[I,J,idx]

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

In [57]: A
Out[57]: 
array([[[5, 1, 5],
        [9, 9, 4]],

       [[0, 9, 8],
        [8, 6, 8]]])

In [59]: idx = np.array([[1,0],[0,2]])

In [60]: m,n = idx.shape

In [61]: I,J = np.ogrid[:m,:n]

In [62]: A[I,J,idx]
Out[62]: 
array([[1, 9],
       [0, 8]])