如何获得 numpy 3-D 数组的对角线分量？

Question

我有一个像

这样的 3-D 数组

a = np.array(
    [[[1, 2],
      [3, 4]],

     [[5, 6],
      [7, 8]]]
)

我想从这个数组中获取对角线分量，我的意思是，

>> np.array([a[i,i,:] for i in range(min(a.shape[0], a.shape[1]))])
array([[1, 2],
       [7, 8]])

有没有更快的方法来完成与上述相同的操作？

Answer 1

np.diagonal 适用于二维以上的数组，您可以使用参数 axis1 和 axis2 指定要查看的轴，对于您的情况，默认值 (axis1 =0, axis2=1) 有效：

np.diagonal(a, axis1=0, axis2=1).T
# array([[1, 2],
#        [7, 8]])

它也适用于在查看的两个轴上具有不同大小的数组：

a = np.array(
    [[[1, 2],
      [3, 4]],

     [[5, 6],
      [7, 8]],

     [[9, 10],
      [11, 12]]]
)

np.diagonal(a).T
#array([[1, 2],
#       [7, 8]])

a = np.array(
    [[[1, 2],
      [3, 4],
      [9, 10]],

     [[5, 6],
      [7, 8],
      [11, 12]]]
)

np.diagonal(a).T
#array([[1, 2],
#       [7, 8]])

或者您可以使用 advanced indexing:

创建范围索引：

idx = np.arange(min(a.shape[:2]))

idx
# array([0, 1])

第一维和第二维使用整型数组索引，第三维使用切片索引：

a[idx, idx, :]
#array([[1, 2],
#       [7, 8]])

Answer 2

代替np.diagonal你也可以使用np.einsum:

>>> a = np.arange(1,9).reshape(2,2,2)
>>> d = np.einsum('iij->ij', a)
>>> d
array([[1, 2],
       [7, 8]])

优点：returns可写视图，可用于直接操作原数组中的对角线：

>>> d[0,0] = -3
>>> a
array([[[-3,  2],
        [ 3,  4]],

       [[ 5,  6],
        [ 7,  8]]])

如何获得 numpy 3-D 数组的对角线分量？

how to get the diagonal components of numpy 3-D array?

numpy

scipy

python-3.x

diagonal