Вычислить ненулевые записи матрицы быстрее в Python

Question

Я хотел бы сделать три вектора из матрицы, суммируя ее ненулевые значения.Вектор значений, вектор индексов строк и вектор индексов столбцов.

Например, если W = [[0. 2. 0.], [0. 10. 0.], [00. 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)

Answer 1

Использование 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]))