从矩阵中提取对角线的Numpy值

Question

我的问题和这个帖子有点相似（这是扩展版）：Numpy 从矩阵中提取行、列和数值。在那个帖子中，我从输入的矩阵中提取了大于零的元素，现在我想提取对角线上的元素。所以在这种情况下，

from numpy import *
import numpy as np

m=np.array([[0,2,4],[4,0,0],[5,4,0]])
dist=[]
index_row=[]
index_col=[]
indices=np.where(matrix>0)
index_col, index_row = indices
dist=matrix[indices]
return index_row, index_col, dist

我们可以得到，

index_row = [1 2 0 0 1]
index_col = [0 0 1 2 2]
dist = [2 4 4 5 4]

现在这就是我想要的，

index_row = [0 1 2 0 1 0 1 2]
index_col = [0 0 0 1 1 2 2 2]
dist = [0 2 4 4 0 5 4 0]

我尝试把原代码的第8行改成这样，

indices=np.where(matrix>0 & matrix.diagonal)

但是出现了这个错误，

我该如何得到我想要的结果呢？请给我一些建议，谢谢！

Answer 1

你可以使用以下方法：

1. 获取一个掩码数组
2. 把掩码的对角线填充为真（True）
3. 选择掩码中为真的元素

下面是代码：

m=np.array([[0,2,4],[4,0,0],[5,4,0]])
mask = m > 0
np.fill_diagonal(mask, True)

m[mask]