Numpy extract row, column and value from a matrix

Question

I have a matrix, and I want to write a script to extract values which are bigger than zero, its row number and column number(because the value belongs to that (row, column)), and here's an example,

from numpy import *
import numpy as np

m=np.array([[0,2,4],[4,0,4],[5,4,0]])
index_row=[]
index_col=[]
dist=[]

I want to store the row number in index_row, the column number in index_col, and the value in dist. So in this case,

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

How to add the codes to achieve this goal? Thanks for giving me suggestions.

Answer 1

You can use numpy.where for this:

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

Answer 2

Though this has been answered already, I often find np.where to be somewhat cumbersome-- though like all things, depends on the circumstance. For this, I'd probably use a zip and a list comprehension:

index_row = [0, 0, 1, 1, 2, 2]
index_col = [1, 2, 0, 2, 0, 1]
zipped = zip(index_row, index_col)
dist = [m[z] for z in zipped]

The zip will give you an iteratable of tuples, which can be used to index numpy arrays.