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

Question

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

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

我可以一次完成吗？

Answer 1

我相信这应该可以解决问题：

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 2

我认为您可以使用理解 - 取决于您将哪些维度定义为 'layer'，等等，但这会起作用：

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

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

Answer 3

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

Selectively access rows of numpy 3D matrix