Python高效数据处理：基于相似度的DataFrame分组聚合优化问询

你的问题本质是基于相似度阈值的连通分量分组——满足阈值的字符串属于同一连通组，且具备传递性。直接两两循环的时间复杂度是O(n²)，对10000行来说完全不可行。最优方案是结合「候选配对筛选」+「并查集（Union-Find）」来实现，把时间复杂度降到可接受范围。

具体实现步骤 #

1. 筛选高相似度候选对（避免全量计算） #

先减少需要计算相似度的配对数，避免1亿次的无效计算。可以用n-gram倒排索引快速找出可能符合阈值的候选：

• 将每个字符串拆成n-gram（比如2-gram或3-gram），例如"apple"拆成{"ap","pp","pl","le"}
• 建立倒排索引：key为n-gram，value为包含该n-gram的字符串索引列表
• 对每个字符串，找出所有共享至少k个n-gram的其他字符串作为候选，仅对这些候选计算真实相似度

代码示例：

from collections import defaultdict

def generate_ngrams(s, n=2):
    return set([s[i:i+n] for i in range(len(s)-n+1)])

# 构建倒排索引
ngram_index = defaultdict(list)
for idx, s in enumerate(df['column1']):
    ngrams = generate_ngrams(s)
    for gram in ngrams:
        ngram_index[gram].append(idx)

# 生成去重的候选配对
candidate_pairs = set()
for idx, s in enumerate(df['column1']):
    ngrams = generate_ngrams(s)
    candidates = set()
    for gram in ngrams:
        candidates.update(ngram_index[gram])
    # 只保留比当前索引大的候选，避免重复计算
    for cand_idx in candidates:
        if cand_idx > idx:
            candidate_pairs.add((idx, cand_idx))

2. 用并查集管理分组 #

拿到候选对后，对每对计算你的相似度函数，若≥0.5则合并它们的组。并查集的合并和查询操作近似O(1)，效率极高。

并查集实现代码：

class UnionFind:
    def __init__(self, size):
        self.parent = list(range(size))
        self.rank = [0]*size
    
    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

# 初始化并查集
uf = UnionFind(len(df))

# 遍历候选对，合并符合条件的组
for idx1, idx2 in candidate_pairs:
    s1 = df['column1'].iloc[idx1]
    s2 = df['column1'].iloc[idx2]
    if your_similarity_function(s1, s2) >= 0.5:  # 替换为你已实现的相似度函数
        uf.union(idx1, idx2)

3. 生成分组ID并聚合 #

给每行分配对应的组ID，再按组聚合求和：

# 生成每个元素的组根节点
df['group_root'] = [uf.find(idx) for idx in range(len(df))]
# 将根节点映射为易读的分组名称
group_mapping = {root: f'group{i+1}' for i, root in enumerate(df['group_root'].unique())}
df['column3'] = df['group_root'].map(group_mapping)

# 按分组聚合column2求和
agg_result = df.groupby('column3')['column2'].sum().reset_index()

额外优化点 #

• 若你的相似度函数支持向量化（比如用TfidfVectorizer转成向量后计算余弦相似度），可批量计算候选对的相似度，进一步提速。
• 调整n-gram的n值和候选筛选条件：n越小，候选越多、召回率越高但计算量越大；n越大则相反，可根据字符串长度和相似度逻辑调整。
• 先过滤长度差异过大的配对（比如长度差超过50%的直接跳过），减少无效计算。

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