在Python中更快地计算矩阵的非零条目

时间:2018-09-08 08:33:32

标签: python matrix

我想从一个矩阵中得出三个向量,总结其非零值。值向量,行索引向量和列索引向量。

例如,如果W = [[0. 2. 0。],[0. 10. 0。],[0. 0. 5.]]。 我希望函数返回([2.0、10.0、5.0],[0、1、2],[1、1、2])。

下面的代码可以完成任务,但是对于大型矩阵来说太慢了。我正在100000的数量级上工作。而且我不知道哪些索引非零。有没有办法加快速度?

from __future__ import division

import numpy as np
import collections
from numpy import *
import copy
#import timing



def nonZeroIndexes(W):
    s = W.shape
    nRows = s[0]
    nColumns = s[1]

    values = []
    row_indexes = []
    column_indexes = []

    for r in xrange(nRows):
        for c in xrange(nColumns):
            if W[r,c] != 0:
                values.append(W[r,c])
                row_indexes.append(r)
                column_indexes.append(c)
    return values, row_indexes, column_indexes

n = 3
W = np.zeros((n,n))

W[0,1] = 2
W[1,1] = 10
W[2,2] = 5

vecs = nonZeroIndexes(W)

1 个答案:

答案 0 :(得分:3)

使用np.nonzero

>>> import numpy as np
>>> W = np.array([[0, 2, 0], [0, 10, 0], [0, 0, 5]])
>>> 
>>> def nonZeroIndexes(W):
...     zero_pos = np.nonzero(W)
...     return (W[zero_pos],) + zero_pos
... 
>>> 
>>> nonZeroIndexes(W)
(array([ 2, 10,  5]), array([0, 1, 2]), array([1, 1, 2]))