如何按列高效遍历scipy.sparse.coo_matrix并生成列级位集？

刚好之前处理过类似的大规模稀疏矩阵列操作需求，针对你100K×1M的场景，直接转成CSC格式就能完美解决按列遍历的问题，而且效率比手动处理COO要高得多。

为什么CSC是最优选择 #

scipy的CSC（Compressed Sparse Column）矩阵本身就是按列存储的，内部用三个核心数组组织数据：

• indptr：长度为列数+1，其中indptr[j]到indptr[j+1]的区间，对应第j列所有非零元素在indices数组中的位置
• indices：存储所有非零元素的行索引
• data：存储非零元素的值（构建位集时用不到这个）

这种结构天生适合按列批量提取数据，完全不需要逐个元素遍历再分配到不同位集里。

具体实现步骤 #

先把你的COO矩阵转成CSC，然后遍历每一列，直接批量获取该列的所有非零行索引，一次性构建对应位集。这里推荐用bitarray库存储位集，它的内存效率极高（100K行的话，每个位集仅占~12KB），非常适配你的大规模场景。

代码示例：

import scipy.sparse as sp
from bitarray import bitarray

# 假设你的原始COO矩阵是coo_mat
csc_mat = coo_mat.tocsc()
num_rows = csc_mat.shape[0]
col_bitsets = []

# 遍历每一列构建位集
for col_idx in range(csc_mat.shape[1]):
    # 获取当前列的所有非零行索引区间
    start = csc_mat.indptr[col_idx]
    end = csc_mat.indptr[col_idx + 1]
    row_indices = csc_mat.indices[start:end]
    
    # 初始化位集并将对应行的位设为1
    col_bit = bitarray(num_rows)
    col_bit.setall(0)
    col_bit[row_indices] = 1
    
    col_bitsets.append(col_bit)

额外优化建议 #

如果你的矩阵非零元素极多，Python循环遍历1M列可能有点耗时，可以考虑用numba给循环加@numba.jit装饰器加速，或者用向量化方式批量处理多列。不过对于大多数场景，上面的基础代码已经足够高效。

另外，若不想用bitarray，也可以用numpy的uint64数组手动实现位集（把每64行打包成一个整数），但bitarray的API更直观，且做了底层优化，没必要重复造轮子。

为什么不直接用COO处理 #

COO矩阵的row、col、data是三个平行数组，默认按行排序。如果要按列遍历，你需要先对col数组排序再分组提取行索引，但排序大规模非零元素数组的开销极大，远不如直接转CSC高效——毕竟scipy的CSC转换是C实现的内部优化操作，速度快很多。

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