如何高效实现按行非零元素个数归一化SciPy的csr_matrix

嘿，这个需求我熟得很！要给SciPy的csr_matrix做按行归一化——让每行元素除以该行非零元素的个数，最简洁高效的做法肯定得利用稀疏矩阵的特性，千万别傻乎乎转成稠密矩阵（大矩阵这么搞内存直接炸）。

最简洁的非原地实现方法 #

这个方法既不会修改原矩阵，又能高效完成计算，核心是利用稀疏矩阵的multiply方法和广播特性：

import numpy as np
from scipy.sparse import csr_matrix

# 先搞个示例矩阵
A = csr_matrix([[1, 2, 0], [0, 3, 4], [5, 0, 0]])

# 第一步：统计每行的非零元素个数
row_nnz = A.getnnz(axis=1)
# 处理全零行：避免除以0，这里把全零行的除数设为1，保证该行元素不变
row_nnz[row_nnz == 0] = 1

# 第二步：执行归一化——每行元素除以对应的非零元素个数
normalized_A = A.multiply(1 / row_nnz[:, np.newaxis])

为什么这个方法高效？

• getnnz(axis=1)是csr_matrix的内置方法，直接基于其行存储结构统计非零元素，速度极快；
• multiply操作只会针对矩阵中的非零元素进行计算，完全跳过零元素，计算量和内存占用都远低于稠密矩阵操作；
• 利用numpy的广播特性，把一维的row_nnz转成列向量后，能自动匹配到每行进行除法运算。

极致内存效率的原地修改方法 #

如果你的矩阵特别大，想省内存，并且允许修改原矩阵，可以直接操作csr_matrix的data属性（存储所有非零元素的数组）：

row_nnz = A.getnnz(axis=1)
# 生成每个非零元素对应的除数：把每行的除数重复对应次数
divisors = np.repeat(1 / row_nnz[row_nnz != 0], row_nnz[row_nnz != 0])
# 直接修改原矩阵的非零元素
A.data /= divisors

这种方法完全不需要额外创建新矩阵，所有操作都在原矩阵的非零元素数组上进行，内存占用最小。

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