CODEWITHSUNDEEP
pythonarraysnumpyindexingdiagonal
Heinz
Heinz

Reputation: 2467

Numpy extract values on the diagonal from a matrix

My question is similar(the expanded version) to this post:Numpy extract row, column and value from a matrix. In that post, I extract elements which are bigger than zero from the input matrix, now I want to extract elements on the diagonal, too. So in this case,

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

we could get,

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

and now this is what I want,

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]

I tried to edit line 8 in the original code to this,

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

but got this error,

enter image description here

How to get the result I want? Please give me some suggestions, thanks!

Upvotes: 0

Views: 2642

Answers (1)

HYRY
HYRY

Reputation: 97331

You can use following method:

  1. get the mask array
  2. fill diagonal of the mask to True
  3. select elements where elements in mask is True

Here is the code:

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

m[mask]

Upvotes: 1

Related Questions