不使用for循环返回nxn矩阵的diag元素

Question

想象一下有一个矩阵数组（nxkxk），我如何在保持原始形状而不使用for循环的情况下返回对角线条目。

例如，在不保持原始形状的情况下，我们可以这样做

np.diagonal(array_of_matrices, axis1=1, axis2=2)

显然我可以这样做，然后重建原始形状，但我想知道是否有更干净的方法。

我什么也没尝试，我已经没有想法了。

np.diag

 不接受 

axis

 参数。

Answer 1

您可以使用广播将

eye

 矩阵乘以 

k x n x n

 形状的数组，以仅返回对角元素。

import numpy as np

n = 5
k = 3
arr = np.arange(k * n * n).reshape(k, n, n)
diags = arr * np.eye(n)

Answer 2

假设：

n = 2 ; k = 3
a = np.arange(1, n*k*k+1).reshape(n,k,k)

array([[[ 1,  2,  3],
        [ 4,  5,  6],
        [ 7,  8,  9]],

       [[10, 11, 12],
        [13, 14, 15],
        [16, 17, 18]]])

您可以使用掩码来识别要掩码的元素：

idx = np.arange(k)

a[:, idx[:,None]!=idx] = 0

输出：

array([[[ 1,  0,  0],
        [ 0,  5,  0],
        [ 0,  0,  9]],

       [[10,  0,  0],
        [ 0, 14,  0],
        [ 0,  0, 18]]])