Scipy CSR稀疏矩阵高效行迭代指南：直接利用内部结构优化性能

import scipy.sparse
from tqdm import tqdm # 假设使用tqdm进行进度显示

# 'matrix' 是一个 scipy.sparse.csr_matrix
# 示例矩阵
matrix = scipy.sparse.random(1000, 1000, format='csr', density=0.01)

for index in tqdm(range(matrix.shape[0]), desc="处理行", leave=False):
    row = matrix.getrow(index)
    values = row.data
    indices = row.indices
    # 对当前行的非零元素进行进一步处理
    # print(f"Row {index}: Indices={indices}, Values={values}")

import scipy.sparse

def iterate_csr_rows_efficiently(matrix, func):
    """
    高效迭代Scipy CSR稀疏矩阵的每一行，并对非零元素及其索引调用指定函数。

    参数:
    matrix (scipy.sparse.csr_matrix): 要迭代的CSR稀疏矩阵。
    func (callable): 一个函数，接受两个参数 (indices, values)，
                     分别代表当前行的非零元素的列索引和值。
    """
    num_rows = matrix.shape[0]
    for row_idx in range(num_rows):
        # 使用indptr获取当前行在data和indices数组中的起始和结束位置
        start_idx = matrix.indptr[row_idx]
        end_idx = matrix.indptr[row_idx + 1]

        # 直接切片获取当前行的非零值和列索引
        # 注意：这里是数组视图，没有进行数据复制
        values = matrix.data[start_idx:end_idx]
        indices = matrix.indices[start_idx:end_idx]

        # 调用用户提供的函数处理当前行数据
        func(indices, values)

# 示例使用
if __name__ == "__main__":
    # 创建一个示例CSR矩阵
    sparse_matrix = scipy.sparse.random(10, 5, format='csr', density=0.5)
    print("原始稀疏矩阵:")
    print(sparse_matrix.toarray())

    def process_row_data(indices, values):
        """一个示例函数，用于处理每一行的非零元素。"""
        if len(indices) > 0:
            print(f"  非零元素索引: {indices}, 对应值: {values}")
        else:
            print("  空行 (无非零元素)")

    print("\n高效迭代结果:")
    iterate_csr_rows_efficiently(sparse_matrix, process_row_data)

import scipy.sparse
import timeit

# 1. 创建测试矩阵
matrix = scipy.sparse.random(10000, 5000, format='csr', density=0.01, random_state=42)

# 2. 定义一个空操作函数
def donothing(*args):
    pass

# 3. 定义三种迭代方法

# 3.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)

# 3.2. COO 转换迭代方法
def get_matrix_rows_coo(matrix, func):
    coo_matrix = matrix.tocoo() # 包含COO转换的时间
    old_i = None
    indices = []
    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(indices, values) # 处理上一行的累积数据
            indices = [j]
            values = [v]
        else: # 同一行
            indices.append(j)
            values.append(v)
        old_i = i

    # 处理最后一行的累积数据
    if indices and values:
        func(indices, values)

# 3.3. 直接CSR方法 (本文推荐)
def get_matrix_rows_efficient_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)

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

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

print("Efficient CSR method:")
%timeit get_matrix_rows_efficient_csr(matrix, donothing)

.getrow() method:
634 ms ± 16.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
COO and iterate method:
270 ms ± 4.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Efficient CSR method:
12.4 ms ± 112 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)