如何在Python中高效计算n-gram重叠矩阵？

问题背景 #

拥有大型文档语料库，需将存在显著bigram重叠的文档合并。示例语料库如下：

corpus = [{'example', 'bigram'}, {'another', 'example'}, {'some', 'outlier'}]

现有实现通过双重循环计算所有文档对的相似度，生成MxM的相似度矩阵，但当文档数量M增大时，O(M²)的时间复杂度导致计算极慢、成本极高。

现有实现代码 #

相似度矩阵计算 #

import numpy as np

sim_matrix = np.zeros((len(corpus), len(corpus)))
for i in range(len(corpus)):
    for j in range(i+1, len(corpus)):
        sim_matrix[i][j] = get_ngram_overlap(corpus[i], corpus[j])
        
np.fill_diagonal(sim_matrix, 1)

Bigram重叠度计算函数 #

def get_ngram_overlap(ngrams_s1, ngrams_s2):
    if min(len(ngrams_s1), len(ngrams_s2)) == 0:
        return 0
    common_ngrams = ngrams_s1 & ngrams_s2
    return len(common_ngrams)/min(len(ngrams_s1), len(ngrams_s2))

示例输出 #

print(sim_matrix)
>> [[1.  0.5  0.]
   [0.  1.   0.]
   [0.  0.   1.]]

优化方案 #

1. 倒排索引减少候选对比对 #

通过构建bigram到文档ID的倒排索引，仅对共享至少一个bigram的文档对计算相似度，大幅减少需要处理的文档对数量：

from collections import defaultdict
import numpy as np

# 构建倒排索引：bigram -> 对应文档ID列表
inverted_index = defaultdict(list)
for doc_id, ngrams in enumerate(corpus):
    for bg in ngrams:
        inverted_index[bg].append(doc_id)

# 收集每个文档的候选对比文档（去重）
candidates = defaultdict(set)
for bg, doc_ids in inverted_index.items():
    doc_count = len(doc_ids)
    for i in range(doc_count):
        for j in range(i+1, doc_count):
            candidates[doc_ids[i]].add(doc_ids[j])
            candidates[doc_ids[j]].add(doc_ids[i])

# 初始化相似度矩阵（对角线默认1）
sim_matrix = np.eye(len(corpus))
# 仅计算候选文档对的相似度
for i in range(len(corpus)):
    for j in candidates.get(i, set()):
        if j > i:  # 避免重复计算
            overlap = get_ngram_overlap(corpus[i], corpus[j])
            sim_matrix[i][j] = overlap
            sim_matrix[j][i] = overlap

优势：当语料库中大部分文档重叠度低时，候选对数量远小于M²，计算效率显著提升。

2. 向量化+稀疏矩阵运算 #

将文档的bigram集合转换为稀疏向量，利用矩阵乘法快速计算共现次数，替代Python循环：

from sklearn.feature_extraction.text import CountVectorizer
import scipy.sparse as sp
import numpy as np

# 将bigram集合转换为字符串格式，适配CountVectorizer
docs = [' '.join(ngrams) for ngrams in corpus]
# 构建bigram的计数向量器（因输入已是bigram，设ngram_range=(1,1)）
vec = CountVectorizer(ngram_range=(1,1))
# 生成稀疏特征矩阵：形状(M, V)，V为唯一bigram的数量
X = vec.fit_transform(docs)

# 计算文档对的共现bigram数量矩阵
cooccur_matrix = X.dot(X.T).toarray()
# 预处理每个文档的bigram长度
doc_lengths = np.array([len(ngrams) for ngrams in corpus])
# 生成每个文档对的最小长度矩阵（用于分母计算）
min_length_matrix = np.minimum(doc_lengths[:, None], doc_lengths[None, :])

# 计算相似度矩阵，处理除0情况
sim_matrix = np.divide(
    cooccur_matrix, 
    min_length_matrix, 
    out=np.zeros_like(cooccur_matrix, dtype=float), 
    where=min_length_matrix != 0
)
# 对角线设为1
np.fill_diagonal(sim_matrix, 1)

优势：利用numpy/scipy的底层C优化实现，比纯Python循环快数倍，稀疏矩阵可有效节省内存。

3. 近似相似度计算（超大规模语料库） #

若无需精确结果，可使用MinHash+LSH（局部敏感哈希）快速定位近似相似的文档对，适用于百万级以上的超大规模语料库：

from datasketch import MinHash, MinHashLSH
import numpy as np

# 初始化LSH索引
lsh = MinHashLSH(threshold=0.3, num_perm=128)
minhashes = []

# 为每个文档生成MinHash并加入LSH
for doc_id, ngrams in enumerate(corpus):
    m = MinHash(num_perm=128)
    for bg in ngrams:
        m.update(bg.encode('utf8'))
    minhashes.append(m)
    lsh.insert(doc_id, m)

# 初始化相似度矩阵
sim_matrix = np.eye(len(corpus))
# 查询每个文档的近似相似文档并计算相似度
for doc_id in range(len(corpus)):
    candidates = lsh.query(minhashes[doc_id])
    for candidate_id in candidates:
        if candidate_id > doc_id:
            overlap = get_ngram_overlap(corpus[doc_id], corpus[candidate_id])
            sim_matrix[doc_id][candidate_id] = overlap
            sim_matrix[candidate_id][doc_id] = overlap

优势：时间复杂度接近线性，能处理超大规模语料库，适合对精度要求不高的场景。

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