如何快速实现Numpy二维字符串数组自匹配？优化大规模计算方案

问题定义 #

给定长度为n的二维字符串numpy数组，需生成n×n的结果数组，规则如下：

• 若两个子数组的元素集合完全相同，结果为2
• 若两个子数组存在公共元素，结果为1
• 无公共元素且集合不同，结果为0

现有实现 #

当前采用逐对遍历的方式实现，核心匹配函数及调用代码如下：

def match(a, b):
    a_set = set(a)
    b_set = set(b)
    if a_set == b_set:
        return 2
    elif a_set & b_set:
        return 1
    else:
        return 0

# 调用方式（无法处理大规模数据）
import numpy as np
import itertools as it
result = np.reshape(np.array([match(i, j) for i, j in it.product(arr, repeat=2)]), (len(arr), len(arr)))

示例 #

输入数组：

arr = np.array([['a', 'b'], ['a', 'b'], ['a', 'c'], ['d', 'e']])

预期输出：

[[2 2 1 0]
 [2 2 1 0]
 [1 1 2 0]
 [0 0 0 2]]

优化需求 #

现有方案时间复杂度为O(n²)，对于n=50000的场景完全无法在合理时间内完成，需要能在数秒内处理该规模数据的优化方案。

优化实现思路 #

通过向量化操作和稀疏矩阵运算替代逐对遍历，将时间复杂度降低至近似O(n*k + n^α)（α<2，取决于数据稀疏度），具体步骤如下：

1. 快速识别集合完全相同的情况 #

将每个子数组的元素集合映射为唯一ID，通过广播比较ID生成结果为2的位置：

import numpy as np
from sklearn.preprocessing import LabelEncoder
from scipy.sparse import csr_matrix

# 给每个子数组的集合分配唯一ID（避免哈希碰撞的稳定实现）
unique_sets = list(set(frozenset(row) for row in arr))
set_to_id = {s: idx for idx, s in enumerate(unique_sets)}
set_ids = np.array([set_to_id[frozenset(row)] for row in arr])

# 生成集合相等的掩码矩阵，对应结果为2
equal_mask = (set_ids[:, None] == set_ids[None, :])
result = np.zeros((len(arr), len(arr)), dtype=int)
result[equal_mask] = 2

2. 快速识别存在公共元素的情况 #

将字符串元素编码为整数，构建稀疏矩阵表示每个子数组包含的元素，通过矩阵乘法快速计算任意两个子数组是否有交集：

# 将所有字符串元素编码为整数
all_elements = arr.flatten()
le = LabelEncoder()
le.fit(all_elements)
int_arr = le.transform(all_elements).reshape(arr.shape)

# 构建稀疏矩阵：行=子数组索引，列=元素编码，值=1表示包含该元素
row_indices = np.repeat(np.arange(len(arr)), arr.shape[1])
col_indices = int_arr.flatten()
sparse_mat = csr_matrix((np.ones(len(row_indices)), (row_indices, col_indices)), 
                        shape=(len(arr), len(le.classes_)))

# 计算交集矩阵：点积>0表示两个子数组有公共元素
intersect_mask = (sparse_mat @ sparse_mat.T).astype(bool)

# 将有交集但集合不同的位置设为1
result[intersect_mask & ~equal_mask] = 1

方案优势 #

• 稀疏矩阵乘法由scipy高度优化，针对高稀疏度数据（如子数组长度远小于唯一元素数）效率极高
• 所有操作均为向量化或批量处理，避免了Python级别的循环开销
• 针对50k规模的输入，该方案可在数秒内完成计算

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