使用 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