如何在稀疏csr矩阵中按行保留Top-K值并掩码其余值？

问题描述 #

给定形状为(M,N)的scipy.sparse.csr_matrix（其中M=50000，N=700000，尺寸较大），需要生成同形状的新稀疏矩阵：将每行中小于该行Top-K最大值的元素置为0，仅保留每行最大的K个值。

示例 #

输入矩阵生成代码：

import scipy.sparse as sparse

R = sparse.random(5,5, density=0.85, format='csr', random_state=42)
print(R.toarray())

输出的稠密矩阵：

array([[0.43194502, 0.78517596, 0.        , 0.59241457, 0.17052412],
       [0.36636184, 0.        , 0.61185289, 0.52475643, 0.94888554],
       [0.19967378, 0.30461377, 0.51423444, 0.06505159, 0.45606998],
       [0.04645041, 0.30424224, 0.29214465, 0.96563203, 0.29122914],
       [0.60754485, 0.13949386, 0.        , 0.        , 0.80839735]])

当K=2时，期望输出的稠密矩阵：

array([[0.        , 0.78517596, 0.        , 0.59241457, 0.        ],
       [0.        , 0.        , 0.61185289, 0.        , 0.94888554],
       [0.        , 0.30461377, 0.51423444, 0.        , 0.        ],
       [0.        , 0.30424224, 0.        , 0.96563203, 0.        ],
       [0.60754485, 0.        , 0.        , 0.        , 0.80839735]])

解决方案 #

针对大尺寸CSR矩阵，直接操作其内部数组（data、indices、indptr）实现高效处理，避免转换为稠密矩阵导致内存溢出。

实现代码 #

import scipy.sparse as sparse
import numpy as np

def csr_row_top_k(matrix, k):
    # 复制原矩阵，避免修改原始数据
    mat = matrix.copy()
    n_rows = mat.shape[0]
    
    for i in range(n_rows):
        # 定位当前行在data和indices中的起止索引
        start_idx = mat.indptr[i]
        end_idx = mat.indptr[i+1]
        row_data = mat.data[start_idx:end_idx]
        
        if len(row_data) <= k:
            # 该行非零元素数量不超过k，无需处理
            continue
        
        # 高效获取该行第k大值的阈值（np.partition比全排序更快）
        threshold = np.partition(row_data, -k)[-k]
        # 生成掩码，保留大于等于阈值的元素（处理多个值等于阈值的情况）
        mask = row_data >= threshold
        # 更新当前行的data和indices
        mat.data[start_idx:end_idx] = row_data[mask]
        mat.indices[start_idx:end_idx] = mat.indices[start_idx:end_idx][mask]
        # 调整indptr数组，匹配修改后的元素数量
        mat.indptr[i+1:] -= (len(row_data) - np.sum(mask))
    
    # 压缩矩阵，清理空元素空间
    mat.eliminate_zeros()
    return mat

代码说明 #

1. 矩阵复制：避免修改输入的原始矩阵。
2. 行定位：通过indptr数组快速定位每行在data和indices中的位置，无需遍历整个矩阵。
3. 高效阈值计算：使用np.partition以O(n)时间复杂度获取Top-K阈值，比全排序更适合大尺寸数据。
4. 元素过滤：通过掩码保留符合条件的元素，确保每行至少保留K个非零值。
5. 结构更新：调整indptr数组并压缩矩阵，优化存储空间和后续操作效率。

测试示例 #

# 生成示例矩阵
R = sparse.random(5,5, density=0.85, format='csr', random_state=42)
# 保留每行Top-2的值
R_top2 = csr_row_top_k(R, k=2)
# 输出结果
print(R_top2.toarray())

输出结果与期望一致。

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