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

时间:2014-06-19 01:54:09

标签: python arrays numpy indexing diagonal

我的问题与此帖子类似(扩展版本):Numpy extract row, column and value from a matrix。在那篇文章中,我从输入矩阵中提取大于零的元素,现在我想在对角线上提取元素。所以在这种情况下,

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)

但是出现了这个错误,

enter image description here

如何获得我想要的结果?请给我一些建议,谢谢!

1 个答案:

答案 0 :(得分:1)

您可以使用以下方法:

  1. 获取掩码数组
  2. 将面具的对角线填充为True
  3. 选择掩码中元素为True的元素
  4. 以下是代码:

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