python - 有选择地访问 numpy 3D 矩阵的行

Question

有没有办法索引一个 3D numpy 矩阵来选择性地抓取每个第 i 层的第 i 行？

例如，我有一个 RxNxR 矩阵，我想抓取第 1 层的第 1 行、第 2 层的第 2 行、第 3 层的第 3 行等，最后得到一个 RxN 矩阵。

我可以在一次操作中做到这一点吗？

标签： pythonnumpymatrixmatrix-indexing

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]])