使用itertools.izip優化迭代

Question

我必須通過所有矩陣的非零值的迭代它看起來像這樣的算法：

for row, col, val in itertools.izip(matrix.row, matrix.col, matrix.data): 
    dostuff(row, col, val) 

我意識到這是通過迭代稀疏的最快方法矩陣numpy，正如在Iterating through a scipy.sparse vector (or matrix)中討論的那樣。

我的問題是，我在每個評估執行函數需要另一個載體，姑且稱之爲vec，和什麼也不做，如果vec[row]等於0，在某些情況下是爲廣大行也是如此。

因此，我不想遍歷矩陣中的所有非零三元組(row, col, val)，其中vec[row] != 0。

我現在要做的就是

import numpy as np 
import scipy.sparse as sp 
import itertools 

N = 10000 
matrix = sp.rand(N, N, density=0.0001, format='coo', dtype=None, random_state=None) 
vec = np.zeroes(N) 
s = 0 
for row, col, val in itertools.izip(matrix.row, matrix.col, matrix.data): 
    if vec[row] != 0: 
     s += vec[row] * val # in reality, some other function is here 

其中工程，如果有不符合vec[row]!=0許多行的工作比原來更快的代碼的簡單和愚蠢的方案。但是，代碼中時vec所有值都爲零的情況下運行速度慢，這是我不能忽略（例如，如果vec=np.ones(len(matrix.data))的情況。

因此，我需要某種izip延長其讓我爲「有條件」通過其輸出迭代，這樣我就喜歡寫東西

for row, col, val in itertools.izip(matrix.row, matrix.col, matrix.data, lambda x: vec[x[0]] !> 0): 
    dostuff(row, col, val) 

你有什麼建議嗎？會是怎樣做到這一點的最快方法是什麼？

Answer 1

你可以只使用numpy的的對行，列和數據進行特殊索引：

which = vec[matrix.row] != 0 

rows = matrix.row[which] 
cols = matrix.col[which] 
data = matrix.data[which] 

for row, col, val in itertools.izip(rows, cols, data): 
    s += vec[row] * val