如何按分区对Numpy二维数组行求和及稀疏矩阵实现方法

嘿，这个问题我之前处理大稀疏矩阵时也纠结过！既然你已经搞懂了稠密数组用np.add.reduceat的分区行求和，那稀疏矩阵的实现其实也有高效的方案，全程不用转成稠密数组浪费内存，完美适配大规模数据场景。

核心思路：利用稀疏矩阵乘法实现分组求和 #

对于Scipy的稀疏矩阵（尤其是最常用的csr_matrix格式），我们可以构造一个稀疏指示矩阵，用它左乘原矩阵，就能快速得到每个分区的行求和结果。这个指示矩阵的作用是：每一行对应一个分区，该行中属于对应分区的原矩阵行位置为1，其余为0。左乘后，原矩阵中同一分区的行就会被自动求和。

具体实现代码 #

import numpy as np
from scipy.sparse import csr_matrix

def sparse_row_reduceat(matrix, partition):
    n_rows = matrix.shape[0]
    # 确保分区包含矩阵的最后一行索引（处理用户可能漏写的情况）
    if partition[-1] != n_rows:
        partition = np.append(partition, n_rows)
    
    # 生成每个原矩阵行对应的分组ID
    group_sizes = np.diff(partition)
    group_ids = np.repeat(np.arange(len(group_sizes)), group_sizes)
    
    # 构造稀疏指示矩阵G：形状为(分组数, 原矩阵行数)
    data = np.ones(n_rows, dtype=matrix.dtype)
    indices = np.arange(n_rows)
    indptr = np.cumsum(np.concatenate([[0], group_sizes]))
    group_matrix = csr_matrix((data, indices, indptr), shape=(len(group_sizes), n_rows))
    
    # 左乘得到分区求和结果
    return group_matrix @ matrix

验证你的示例 #

用你给出的测试数据跑一遍：

# 构造原稀疏矩阵
test_matrix = csr_matrix([[0,1,1,1], [2,0,1,1], [0,0,0,1], [5,1,0,0]])
# 给定分区
partition = [0, 1]
# 计算结果
result = sparse_row_reduceat(test_matrix, partition)
# 转成数组查看
print(result.toarray())

输出完全符合预期：

[[0 1 1 1]
 [7 1 1 2]]

额外说明 #

• 这个方法全程在稀疏域操作，内存占用极低，哪怕是百万级行的稀疏矩阵也能轻松处理；
• 如果你的稀疏矩阵是csc_matrix或其他格式，建议先转成csr_matrix（用.tocsr()方法），因为CSR格式对行操作的支持更高效；
• 分区不需要手动补全到矩阵行数，函数里已经做了兼容处理。

内容的提问来源于stack exchange，提问作者ffffffyyyy