高效迭代 SciPy CSR 稀疏矩阵中每行的非零元素

import scipy.sparse
from tqdm import tqdm # 用于进度显示

def get_matrix_original(matrix, func):
    """
    使用 matrix.getrow() 方法迭代每一行的非零元素。
    """
    for index in tqdm(range(matrix.shape[0]), desc="Processing rows", leave=False):
        row = matrix.getrow(index)
        indices = row.indices
        values = row.data
        func(indices, values) # 对当前行的索引和值进行处理

def get_matrix_rows_coo(matrix, func):
    """
    将 CSR 矩阵转换为 COO 格式后迭代非零元素。
    """
    coo_matrix = matrix.tocoo() # 转换到 COO 格式
    old_i = None
    current_indices = []
    current_values = []

    # 遍历 COO 格式的行、列、值
    for i, j, v in zip(coo_matrix.row, coo_matrix.col, coo_matrix.data):
        if i != old_i: # 新的一行开始
            if old_i is not None:
                func(current_indices, current_values) # 处理上一行的非零元素
            current_indices = [j]
            current_values = [v]
        else: # 同一行
            current_indices.append(j)
            current_values.append(v)
        old_i = i

    # 处理最后一行的非零元素
    if current_indices and current_values:
        func(current_indices, current_values)

def get_matrix_rows_optimized_csr(matrix, func):
    """
    直接利用 CSR 矩阵的 indptr 结构高效迭代每一行的非零元素。
    """
    rows = matrix.shape[0]
    for index in range(rows):
        # 使用 indptr 确定当前行的非零元素在 data 和 indices 数组中的范围
        indptr_start = matrix.indptr[index]
        indptr_end = matrix.indptr[index + 1]

        # 直接切片获取当前行的值和列索引
        values = matrix.data[indptr_start:indptr_end]
        indices = matrix.indices[indptr_start:indptr_end]

        func(indices, values) # 对当前行的索引和值进行处理

import scipy.sparse
import numpy as np
import timeit

# 假设一个空操作函数，用于模拟实际处理
def donothing(*args):
    pass

# 模拟一个稀疏矩阵
matrix = scipy.sparse.random(10000, 5000, format='csr', density=0.01, random_state=42)

# --- 1. getrow() 方法 ---
def get_matrix_original(matrix, func):
    for index in range(matrix.shape[0]):
        row = matrix.getrow(index)
        indices = row.indices
        values = row.data
        func(indices, values)

# --- 2. COO 转换和迭代方法 ---
def get_matrix_rows_coo(matrix, func):
    coo_matrix = matrix.tocoo()
    old_i = None
    indices = []
    values = []

    for i, j, v in zip(coo_matrix.row, coo_matrix.col, coo_matrix.data):
        if i != old_i:
            if old_i is not None:
                func(indices, values)
            indices = [j]
            values = [v]
        else:
            indices.append(j)
            values.append(v)
        old_i = i

    if indices and values: # 处理最后一组
        func(indices, values)

# --- 3. 优化后的 CSR 直接访问方法 ---
def get_matrix_rows_optimized_csr(matrix, func):
    rows = matrix.shape[0]
    for index in range(rows):
        indptr_start = matrix.indptr[index]
        indptr_end = matrix.indptr[index + 1]
        values = matrix.data[indptr_start:indptr_end]
        indices = matrix.indices[indptr_start:indptr_end]
        func(indices, values)

# 运行基准测试
print(".getrow() method:")
%timeit get_matrix_original(matrix, donothing)

print("COO conversion and iterate method:")
%timeit get_matrix_rows_coo(matrix, donothing)

print("Optimized CSR method:")
%timeit get_matrix_rows_optimized_csr(matrix, donothing)

.getrow() method
634 ms ± 16.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

COO conversion and iterate method
270 ms ± 4.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Optimized CSR method
12.4 ms ± 112 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)