如何快速按索引对NumPy数组求和及优化稀疏矩阵除法

针对数亿行×百万列、数百万非零值的csr_matrix，直接除法会转为稠密矩阵导致内存溢出，原循环实现效率低下，优化方案基于稀疏矩阵结构实现全向量化操作：

原实现问题 #

遍历非零元素的行/列索引并逐次求和，属于**O(n²)**时间复杂度，百万级非零元素耗时极长。

优化方案：利用稀疏矩阵原生属性 #

1. 行归一化（每行元素除以行和）

利用csr_matrix的indptr属性（行指针，标记每行非零元素在data中的起止位置）快速生成对齐的行和数组：

import scipy.sparse as sp

# 假设csr_arr为目标稀疏矩阵
shape = csr_arr.shape

# 计算每行的和，处理全零行（避免除零，可根据业务调整替换值）
row_sums = csr_arr.sum(axis=1).A.flatten()
row_sums[row_sums == 0] = 1

# 生成与data等长的行和数组：每行的和重复对应行非零元素的次数
row_sums_repeated = np.repeat(row_sums, np.diff(csr_arr.indptr))

# 构建行归一化后的稀疏矩阵
prob_b_given_a_data = csr_arr.data / row_sums_repeated
csr_prob_b_given_a = sp.csr_matrix((prob_b_given_a_data, csr_arr.indices, csr_arr.indptr), shape=shape)

2. 列归一化（每列元素除以列和）

利用csr_matrix的indices属性（非零元素的列索引）直接匹配列和，若矩阵规模极大，转csc_matrix计算列和效率更高：

# 计算每列的和，处理全零列
col_sums = csr_arr.sum(axis=0).A.flatten()
col_sums[col_sums == 0] = 1

# 生成与data等长的列和数组：直接通过列索引匹配对应列的和
col_sums_matched = col_sums[csr_arr.indices]

# 构建列归一化后的稀疏矩阵
prob_a_given_b_data = csr_arr.data / col_sums_matched
csr_prob_a_given_b = sp.csr_matrix((prob_a_given_b_data, csr_arr.indices, csr_arr.indptr), shape=shape)

优化原理 #

• 行求和通过csr_matrix.sum(axis=1)高效完成，无需遍历所有非零元素；
• 利用indptr和indices直接生成与data对齐的求和数组，完全避免循环，时间复杂度降至O(n)；
• 全程基于稀疏矩阵结构操作，不会转为稠密矩阵，彻底解决内存溢出问题。

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