python - 使用 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 的所有值都存在的情况下，代码运行缓慢。是非零的，这是我不能忽略的情况(例如， if 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