如何快速查找两个不同大小SciPy稀疏矩阵的相同行索引？

问题背景 #

现有n×j的SciPy稀疏矩阵A和m×j的SciPy稀疏矩阵B（m<n），需要找出A中所有与B的行完全相同的行的索引。常规的循环遍历、转换为Numpy密集数组的方法因矩阵规模过大无法生效。

示例代码 #

import numpy as np
from scipy.sparse import csc_matrix

# 构建稀疏矩阵A
row = np.array([0, 2, 2, 0, 1, 2])
col = np.array([0, 0, 1, 2, 2, 2])
data = np.array([1, 3, 3, 4, 5, 6])
A = csc_matrix((data, (row, col)), shape=(5, 3))
print("A的密集形式：")
print(A.toarray())
# 输出：
# [[1 0 4]
#  [0 0 5]
#  [3 3 6]
#  [0 0 0]
#  [0 0 0]]

# 构建稀疏矩阵B
row = np.array([0, 2, 2, 0, 1, 2])
col = np.array([0, 0, 1, 2, 2, 2])
data = np.array([1, 2, 3, 4, 5, 6])
B = csc_matrix((data, (row, col)), shape=(4, 3))
print("\nB的密集形式：")
print(B.toarray())
# 输出：
# [[1 0 4]
#  [0 0 5]
#  [2 3 6]
#  [0 0 0]]

期望输出 #

def some_function(A,B): 
    # 实现逻辑
    return indices

>>> some_function(A,B)
[0, 1, 3, 4]

解决方案 #

针对稀疏矩阵的特性，我们可以通过提取每行的非零元素特征（列索引+对应值）生成可哈希的键，利用集合快速查找匹配行，避免转换为密集数组带来的内存压力。

实现代码 #

import numpy as np
from scipy.sparse import csc_matrix

def get_row_key(row):
    """将稀疏矩阵的行转换为可哈希的键"""
    # 获取当前行的非零元素列索引和对应值
    cols = row.indices
    vals = row.data
    # 按列索引排序，确保相同行的键一致
    sorted_pairs = sorted(zip(cols, vals))
    return tuple(sorted_pairs)

def find_matching_rows(A, B):
    # 预处理B，将所有行的键存入集合
    b_row_keys = set()
    for i in range(B.shape[0]):
        b_row = B.getrow(i)
        b_row_keys.add(get_row_key(b_row))
    
    # 遍历A的每一行，检查是否存在匹配
    matching_indices = []
    for idx in range(A.shape[0]):
        a_row = A.getrow(idx)
        if get_row_key(a_row) in b_row_keys:
            matching_indices.append(idx)
    
    return matching_indices

# 测试示例
if __name__ == "__main__":
    # 构建示例矩阵A和B
    row_a = np.array([0, 2, 2, 0, 1, 2])
    col_a = np.array([0, 0, 1, 2, 2, 2])
    data_a = np.array([1, 3, 3, 4, 5, 6])
    A = csc_matrix((data_a, (row_a, col_a)), shape=(5, 3))

    row_b = np.array([0, 2, 2, 0, 1, 2])
    col_b = np.array([0, 0, 1, 2, 2, 2])
    data_b = np.array([1, 2, 3, 4, 5, 6])
    B = csc_matrix((data_b, (row_b, col_b)), shape=(4, 3))

    print(find_matching_rows(A, B))  # 输出: [0, 1, 3, 4]

方案说明 #

1. 稀疏行特征提取：通过get_row_key函数，将每行的非零元素（列索引，值）对排序后转为元组，确保内容相同的行生成完全一致的键（即使非零元素存储顺序不同）。
2. 集合快速查找：先将B的所有行特征存入集合，遍历A时只需O(1)时间判断当前行是否存在匹配，整体时间复杂度为O(n + m)（基于稀疏矩阵非零元素的处理量）。
3. 内存友好：全程仅处理稀疏矩阵的非零元素，无需加载整个密集矩阵，适合大规模稀疏矩阵场景。

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