行向量与列向量的高效逐元素比较及稀疏矩阵实现问询

问题背景 #

给定行向量a = np.array([1, 2, 3])和列向量b = np.array([[1], [2], [3]])，执行c = a==b可实现逐元素比较，返回结果：

>>> c
array([[ True, False, False],
       [False,  True, False],
       [False, False,  True]])

但当元素数量极大时，这种稠密矩阵的比较方式会占用大量内存。现有如下定义的Scipy稀疏矩阵a和b：

from scipy.sparse import csr_matrix
import numpy as np

data = np.array([1, 2, 3])
row = np.array([0, 1, 2])
col = np.array([0, 0, 0])

a = csr_matrix((data, (row, col)), shape=(3, 1))
b = csr_matrix((data, (col, row)), shape=(1, 3))

需要利用这两个稀疏矩阵高效计算得到对应稀疏矩阵c。

解决方案 #

核心思路是只追踪值相等的位置，避免生成全量稠密矩阵，具体步骤如下：

1. 提取稀疏矩阵的非零值
   由于这里的a和b所有元素都是非零的，直接取data属性即可：

   a_vals = a.data
   b_vals = b.data
   

2. 建立值到索引的映射
   用哈希表记录b中每个值对应的列索引，方便快速匹配a中的值：

   from collections import defaultdict
   
   val_to_b_cols = defaultdict(list)
   for col_idx, val in enumerate(b_vals):
       val_to_b_cols[val].append(col_idx)
   

3. 收集相等位置的索引
   遍历a的每一行，找到b中值相等的列，记录对应的行、列索引：

   c_rows = []
   c_cols = []
   c_data = []
   
   for row_idx, val in enumerate(a_vals):
       if val in val_to_b_cols:
           for col_idx in val_to_b_cols[val]:
               c_rows.append(row_idx)
               c_cols.append(col_idx)
               c_data.append(1)  # 稀疏矩阵用1表示True，后续可转为bool类型
   

4. 构建稀疏矩阵c
   用收集到的行、列和数据构建CSR格式的稀疏矩阵：

   c = csr_matrix((c_data, (c_rows, c_cols)), shape=(a.shape[0], b.shape[1]))
   

结果验证 #

将稀疏矩阵转为稠密数组查看，结果与原numpy方式一致：

print(c.toarray().astype(bool))
# 输出：
# [[ True False False]
#  [False  True False]
#  [False False  True]]

这种方式仅存储值相等的非零位置，内存占用远低于稠密矩阵比较，尤其在元素量大、相等位置稀疏的场景下，效率提升显著。

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