如何删除csr_matrix中非零小值并提升稀疏性？优化慢处理问题

嘿，我之前刚好处理过类似规模的稀疏矩阵，直接用cgc[cgc>1]确实慢得离谱——这是因为CSR格式的布尔索引操作在底层会做很多额外的结构校验和转换，尤其是当非零元素数量达到20亿级别的时候。给你两个经过实测的高效方案，比原生索引快好几倍：

方案一：直接操作CSR的内部数组（最快，内存友好） #

CSR矩阵的核心数据都存在三个NumPy数组里：data（所有非零值）、indices（对应列索引）、indptr（每行非零元素的起始/结束位置）。我们可以直接对data做阈值筛选，然后重新构建CSR矩阵，全程用NumPy的向量化操作，避开scipy稀疏矩阵的慢逻辑：

import scipy.sparse as sp
import numpy as np

# 加载矩阵
cgc = sp.load_npz('/root/cg.npz')
print(f"原始非零元素数: {cgc.count_nonzero()}")

# 1. 生成非零值的筛选掩码
mask = cgc.data > 1

# 2. 筛选出符合条件的data和indices
new_data = cgc.data[mask]
new_indices = cgc.indices[mask]

# 3. 重新计算新的indptr（关键：统计每行剩余的非零元素数）
# 用reduceat快速计算每行的有效元素数量
row_valid_counts = np.add.reduceat(mask, cgc.indptr[:-1])
# 累加得到新的indptr
new_indptr = np.concatenate([[0], np.cumsum(row_valid_counts)])

# 4. 构建新的CSR矩阵
cgc_filtered = sp.csr_matrix((new_data, new_indices, new_indptr), shape=cgc.shape)
print(f"筛选后非零元素数: {cgc_filtered.count_nonzero()}")

这个方法的优势是：

• 全程基于NumPy的C级向量化操作，比scipy的稀疏索引快很多
• 不需要转换矩阵格式，内存占用和原矩阵基本一致（只多了几个临时数组）

方案二：转成COO格式处理（更简洁，适合内存充足的情况） #

COO格式的矩阵结构更简单（直接存储所有非零元素的行、列、值），筛选逻辑更直观。如果你的服务器内存足够（能容纳三个20亿级别的数组），可以试试这个：

import scipy.sparse as sp

cgc = sp.load_npz('/root/cg.npz')
print(f"原始非零元素数: {cgc.count_nonzero()}")

# 转成COO格式
coo = cgc.tocoo()
# 筛选符合条件的元素
mask = coo.data > 1
# 构建新的COO矩阵再转回CSR
cgc_filtered = sp.coo_matrix(
    (coo.data[mask], (coo.row[mask], coo.col[mask])),
    shape=cgc.shape
).tocsr()
print(f"筛选后非零元素数: {cgc_filtered.count_nonzero()}")

这个方法代码更简洁，但转换COO格式会额外占用一些内存（需要存储行索引数组），如果你的内存吃紧的话，优先选方案一。

为什么原生的cgc[cgc>1]这么慢？ #

scipy的稀疏矩阵布尔索引cgc[cgc>1]在底层会先创建一个布尔型的稀疏矩阵（cgc>1），然后再做矩阵乘法或者索引匹配，这个过程会涉及大量的稀疏结构校验和临时矩阵创建，对于20亿级别的非零元素来说，这些额外操作会拖慢整个流程。而我们直接操作内部数组的方式，跳过了这些冗余步骤，直接针对有效数据处理，效率自然高很多。

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