如何实现集群内名称相似度评分的高效矩阵计算并避免内存过载？

问题核心分析 #

原方案使用rapidfuzz.process.cdist生成全量相似度矩阵，在百万级数据下会因内存不足崩溃；尝试的全局分块方案错误地将同分组内的相似项拆分到不同块，导致无法识别相似关系，最终ID分配错误。

关键优化方向：

• 限定分组计算：需求明确要求同一cluster_number内识别相似名称，跨组无需比较，这能大幅降低计算规模。
• 避免全量矩阵：用逐项查询+并查集（Union-Find）替代全量矩阵，内存仅与单分组大小相关，而非全局数据量。

实现代码 #

1. 并查集工具类（用于合并相似名称组） #

class UnionFind:
    def __init__(self, elements):
        self.parent = {elem: elem for elem in elements}
        self.rank = {elem: 0 for elem in elements}
    
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def union(self, x, y):
        x_root = self.find(x)
        y_root = self.find(y)
        if x_root == y_root:
            return
        if self.rank[x_root] < self.rank[y_root]:
            self.parent[x_root] = y_root
        else:
            self.parent[y_root] = x_root
            if self.rank[x_root] == self.rank[y_root]:
                self.rank[x_root] += 1

2. 核心处理函数 #

from rapidfuzz import fuzz, process
import pandas as pd

def assign_group_ids(df, similarity_threshold=50):
    name_to_id = {}
    current_global_id = 1

    # 按cluster_number分组处理，严格限定相似性计算范围
    for cluster_id, group_df in df.groupby('cluster_number'):
        group_names = group_df['name'].tolist()
        uf = UnionFind(group_names)

        # 遍历组内每个名称，查找相似项并合并
        for name in group_names:
            if name in name_to_id:
                continue
            # 提取组内相似度≥阈值的结果，排除自身
            similar_results = process.extract(
                name, 
                group_names, 
                scorer=fuzz.ratio, 
                score_cutoff=similarity_threshold
            )
            similar_names = [res[0] for res in similar_results if res[0] != name]
            # 合并所有相似名称到同一集合
            for sim_name in similar_names:
                uf.union(name, sim_name)
        
        # 为当前分组的每个连通分量分配全局唯一ID
        cluster_id_map = {}
        for name in group_names:
            root = uf.find(name)
            if root not in cluster_id_map:
                cluster_id_map[root] = current_global_id
                current_global_id += 1
            name_to_id[name] = cluster_id_map[root]
    
    # 映射ID到原DataFrame
    df['id'] = df['name'].map(name_to_id)
    return df

3. 测试验证 #

d_test = {
    'name' : ['South Beach', 'Dog', 'Bird', 'Ant', 'Big Dog', 'Beach', 'Dear', 'Cat'],
    'cluster_number' : [1, 2, 3, 3, 2, 1, 4, 2]
}
df_test = pd.DataFrame(d_test)
result_df = assign_group_ids(df_test)
print(result_df.sort_values('cluster_number').reset_index(drop=True))

输出结果（与预期一致）：

name  cluster_number  id
0  South Beach               1   1
1        Beach               1   1
2          Dog               2   2
3      Big Dog               2   2
4          Cat               2   3
5          Ant               3   4
6         Bird               3   5
7         Dear               4   6

方案优势 #

1. 内存高效：无需生成全量相似度矩阵，仅针对单分组内的名称逐次查询相似项，内存占用与单分组大小正相关，百万级数据只要分组规模合理（如单组≤10万条）即可稳定运行。
2. 结果准确：通过并查集保证同一分组内所有相似名称被分配同一ID，不会出现分块方案的拆分错误。
3. 可扩展性：支持调整相似度阈值、替换匹配算法（如用fuzz.partial_ratio优化子串匹配场景），适配不同业务需求。

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