python - Numpy extract values on the diagonal from a matrix -

python - Numpy extract values on the diagonal from a matrix -

my question similar(the expanded version) post:numpy extract row, column , value matrix. in post, extract elements bigger 0 input matrix, want extract elements on diagonal, too. in case,

from numpy import * import numpy 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] homecoming index_row, index_col, dist

we get,

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

and 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 edit line 8 in original code this,

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

but got error,

how result want? please give me suggestions, thanks!

you can utilize next method:

get mask array fill diagonal of mask true select elements elements in mask true

here code:

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

python arrays numpy indexing diagonal