Numpy 从矩阵中提取行、列和值

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

您可以为此使用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])

Answer 2

虽然已经回答了这个问题，但我经常发现np.where 有点麻烦——尽管像所有事情一样，视情况而定。为此，我可能会使用zip 和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]

zip 将为您提供一个可迭代的元组，可用于索引 numpy 数组。