numpy - 多维 numpy 数组的对角线

Question

是否有更 pythonic 的方式来执行以下操作:

import numpy as np
def diagonal(A):
    (x,y,y) = A.shape
    diags = []
    for a in A: diags.append(np.diagonal(a))
    result = np.vstack(diags)
    assert result.shape == (x,y)
    return result

Answer 1

方法 #1

一个干净的方式是使用np.diagonal在输入数组的转置版本上，像这样 -

np.diagonal(A.T)

基本上，我们使用 A.T 翻转输入数组的维度，让 np.diagonal 使用最后两个轴来提取对角线元素，因为默认情况下否则它会使用前两个轴。最好的事情是这适用于任意维数的数组。

方法 #2

这是使用 advanced and basic indexing 组合的另一种方法-

m,n = A.shape[:2]
out = A[np.arange(m)[:,None],np.eye(n,dtype=bool)]

还可以通过基本索引 进行一些 reshape -

out = A.reshape(m,-1)[:,np.eye(n,dtype=bool).ravel()]

sample 运行-

In [87]: A
Out[87]: 
array([[[73, 52, 62],
        [20,  7,  7],
        [ 1, 68, 89]],

       [[15, 78, 98],
        [24, 22, 35],
        [19,  1, 91]],

       [[ 5, 37, 64],
        [22,  4, 43],
        [84, 45, 12]],

       [[24, 45, 42],
        [70, 45,  1],
        [ 6, 48, 60]]])

In [88]: np.diagonal(A.T)
Out[88]: 
array([[73,  7, 89],
       [15, 22, 91],
       [ 5,  4, 12],
       [24, 45, 60]])

In [89]: m,n = A.shape[:2]

In [90]: A[np.arange(m)[:,None],np.eye(n,dtype=bool)]
Out[90]: 
array([[73,  7, 89],
       [15, 22, 91],
       [ 5,  4, 12],
       [24, 45, 60]])