如何向量化并加速Pandas DataFrame文本相似度计算的双重循环

问题背景 #

给定如下DataFrame，需要为同一cluster_number下的相似name分配唯一ID：

import pandas as pd

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)

需求示例：

• Cluster 1中的South Beach和Beach相似度高，分配同一ID
• Cluster 2中的Dog和Big Dog归为一组，Cat单独分配ID
• 孤立项（如Cluster3的Bird/Ant、Cluster4的Dear）各自分配唯一ID

原实现的双重iterrows()循环在测试数据上可行，但百万级数据下性能完全无法忍受，必须重构为无循环的高效方案。

优化方案1：结合TheFuzz.Process + 图连通分量（NetworkX） #

核心思路：按Cluster分组，每组内通过批量模糊匹配构建相似关系图，再用连通分量算法给相似组分配ID，彻底规避双重循环。

依赖安装

先装加速依赖（必须装，否则百万级数据速度依然堪忧）：

pip install thefuzz python-Levenshtein networkx

实现代码

import pandas as pd
from thefuzz import process, fuzz
import networkx as nx

def assign_similar_ids(df, similarity_threshold=50):
    df['id'] = 0
    current_unique_id = 1

    # 按cluster分组处理，避免跨组无效计算
    for cluster_id, group in df.groupby('cluster_number'):
        group_names = group['name'].tolist()
        name_to_group_idx = {name: idx for idx, name in enumerate(group_names)}
        # 构建无向图存储相似关系
        similarity_graph = nx.Graph()
        similarity_graph.add_nodes_from(group_names)

        # 批量提取每组内的相似项
        for name in group_names:
            # 只在当前组内匹配，限制返回所有符合阈值的结果
            matched_items = process.extract(name, group_names, limit=None, scorer=fuzz.ratio)
            for matched_name, score in matched_items:
                if score >= similarity_threshold and name != matched_name:
                    similarity_graph.add_edge(name, matched_name)

        # 给每个连通分量分配统一ID
        for connected_group in nx.connected_components(similarity_graph):
            for name in connected_group:
                group_idx = name_to_group_idx[name]
                df.loc[group.index[group_idx], 'id'] = current_unique_id
            current_unique_id += 1

        # 给孤立节点（无相似项）分配单独ID
        isolated_names = [name for name in group_names if nx.degree(similarity_graph, name) == 0]
        for name in isolated_names:
            group_idx = name_to_group_idx[name]
            df.loc[group.index[group_idx], 'id'] = current_unique_id
            current_unique_id += 1

    return df

# 测试运行
df_result = assign_similar_ids(df_test, similarity_threshold=50)
print(df_result)

输出结果

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

优化方案2：向量化 + 并查集（无需NetworkX） #

如果不想引入NetworkX依赖，可以用Numpy批量计算相似度，结合并查集（Union-Find）算法标记连通组：

import pandas as pd
import numpy as np
from thefuzz import fuzz

def batch_fuzzy_ratio(a, b):
    # 批量计算字符串相似度，转为Numpy数组
    return np.array([fuzz.ratio(str(x), str(y)) for x, y in zip(a, b)])

def assign_ids_vectorized(df, similarity_threshold=50):
    df['id'] = 0
    current_unique_id = 1

    for cluster_id, group in df.groupby('cluster_number'):
        names = group['name'].values
        group_size = len(names)
        # 生成上三角索引，避免重复计算两两配对
        idx_i, idx_j = np.triu_indices(group_size, k=1)
        # 批量计算所有配对的相似度
        similarity_scores = batch_fuzzy_ratio(names[idx_i], names[idx_j])
        # 筛选符合阈值的有效配对
        valid_pairs = np.where(similarity_scores >= similarity_threshold)[0]

        # 初始化并查集
        parent = np.arange(group_size)
        
        def find_root(u):
            while parent[u] != u:
                parent[u] = parent[parent[u]]  # 路径压缩优化
                u = parent[u]
            return u
        
        def union_nodes(u, v):
            u_root = find_root(u)
            v_root = find_root(v)
            if u_root != v_root:
                parent[v_root] = u_root

        # 合并所有相似节点
        for pair_idx in valid_pairs:
            u, v = idx_i[pair_idx], idx_j[pair_idx]
            union_nodes(u, v)

        # 给每个连通分量分配ID
        root_to_id = {}
        for idx in range(group_size):
            root = find_root(idx)
            if root not in root_to_id:
                root_to_id[root] = current_unique_id
                current_unique_id += 1
            df.loc[group.index[idx], 'id'] = root_to_id[root]

    return df

# 测试运行
df_result = assign_ids_vectorized(df_test, similarity_threshold=50)
print(df_result)

关键性能优化点 #

1. 必须安装python-Levenshtein：TheFuzz默认用纯Python实现编辑距离计算，速度极慢；安装C扩展的Levenshtein后，计算速度可提升10~100倍。
2. 按Cluster分组处理：只在同一Cluster内计算相似度，大幅减少无效计算量。
3. 连通分量/并查集替代双重循环：将相似关系转化为图连通问题，时间复杂度从O(n²)降至近似O(n)，完全适配百万级数据。
4. 批量处理：利用Numpy或TheFuzz的批量接口，减少Python循环的开销。

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