从矩阵中提取行、列和值

Question

我有一个矩阵，想写个脚本来提取那些大于零的值，以及它们所在的行号和列号（因为这个值是属于那个（行，列）位置的）。下面是一个例子，

from numpy import *
import numpy as np

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

我想把行号存储在index_row里，把列号存储在index_col里，把值存储在dist里。所以在这个例子中，

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

我该怎么加代码来实现这个目标呢？谢谢你们给我的建议。

Answer 1

虽然这个问题已经有人回答过了，但我发现使用 np.where 有时候挺麻烦的——不过这也要看具体情况。对于这个问题，我可能会用到 zip 和 列表推导式：

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]

使用 zip 可以得到一个可迭代的元组，这些元组可以用来索引 numpy 数组。

Answer 2

你可以使用 numpy.where 来实现这个功能：

>>> 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])