Селективный захват рядов матрицы Numpy

Question

Есть ли способ индексировать трехмерную матрицу, чтобы выборочно захватить i-ую строку каждого слоя?

Например, у меня есть матрица RxNxR, и я хочу взять 1-ую строку 1-го уровня, 2-ю строку 2-го уровня, 3-й ряд 3-го уровня и т. д. и получить матрицу RxN.

Могу ли я сделать это за одну операцию?

Answer 1

Используя функцию numpy.diagonal:

a = np.arange(27).reshape(3, 3, 3)
np.diagonal(a, offset=0, axis1=0, axis2=1).T

дает

array([[ 0,  1,  2],
       [12, 13, 14],
       [24, 25, 26]])

Answer 2

Я считаю, что это должно сработать

import numpy as np
matrix = ...
result = np.zeros((matrix.shape[0], matrix.shape[1]))

for i in range(0, len(matrix)):
    result[i] = matrix[i][i]

Answer 3

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

a = np.array([...])

[a[:, i, 0] for i in range(np.shape(a)[1])]