基于余弦相似度分组字符串列表：相似项未正确分组的问题排查

我有一个包含约2000个元素的字符串列表，示例为my_list = ['policeman', 'police officers', 'police force', ...]。需求是基于余弦相似度对列表词汇分组：相似度高于0.7的归为同一组，且每个词汇仅能属于一个分组。

我编写了如下Python函数，但运行后发现policeman与police officers未被正确分组，期望输出示例为[['policeman', 'police officers', 'police force'], ['apple', 'banana','pineapple']]：

def subject_similarity_grouped(subj_list, threshold):
            
    embeddings = {word: nlp(word).vector for word in subj_list}
    
    # cosine similarity and grouping
    # create a list to hold the groups
    groups = []

    # iterate over each word in the list
    for word in subj_list:
        # check if the word is already in a group
        in_a_group = False
        for group in groups:
            if word in group:
                in_a_group = True
                break

        # if the word is not in a group, create a new group for it
        if not in_a_group:
            # create a new group with the current word
            new_group = [word]

            # retrieve the embedding for the current word
            embedding1 = embeddings[word]

            # iterate over the remaining words and add them to the current group if they are similar enough
            for other_word in subj_list:
                # skip the current word
                if other_word == word:
                    continue

                # check if the other word is already in a group
                in_a_group = False
                for group in groups:
                    if other_word in group:
                        in_a_group = True
                        break

                # if the other word is not in a group, retrieve its embedding and compute its similarity to the current word
                if not in_a_group:
                    embedding2 = embeddings[other_word]
                    similarity = cosine_similarity(embedding1.reshape(1, -1), embedding2.reshape(1, -1))[0][0]
                    # if the similarity is above the threshold, add the word to the current group
                    if similarity > threshold:
                        new_group.append(other_word)

            # add the new group to the list of groups
            groups.append(new_group)
            
    # remove if the grouped subject list only has one entity       
    groups = [lst for lst in groups if len(lst) > 1]
    return groups

问题根源 #

1. 单一匹配逻辑缺陷：当前代码仅用新组的第一个词和其他词计算相似度，若policeman与police officers的相似度刚好低于0.7，但police officers和组内后续加入的police force相似度高于0.7，这种间接相关的词会被遗漏。
2. 低效的分组检查：每次判断词是否已分组都遍历所有组，对于2000个元素的列表，时间复杂度会达到O(n²)，运行效率极低。

修复后的代码 #

from sklearn.metrics.pairwise import cosine_similarity
# 假设nlp是已加载的spaCy模型，比如nlp = spacy.load("en_core_web_lg")

def subject_similarity_grouped(subj_list, threshold=0.7):
    # 预计算所有词的嵌入向量
    embeddings = {word: nlp(word).vector for word in subj_list}
    # 用集合跟踪已分组的词，O(1)查询效率
    grouped_words = set()
    groups = []

    for word in subj_list:
        if word in grouped_words:
            continue
        
        # 初始化新组，用队列实现广度优先匹配
        current_group = [word]
        grouped_words.add(word)
        queue = [word]

        while queue:
            current_word = queue.pop(0)
            current_embedding = embeddings[current_word]

            for other_word in subj_list:
                if other_word in grouped_words or other_word == current_word:
                    continue
                
                other_embedding = embeddings[other_word]
                similarity = cosine_similarity(
                    current_embedding.reshape(1, -1), 
                    other_embedding.reshape(1, -1)
                )[0][0]

                if similarity > threshold:
                    current_group.append(other_word)
                    grouped_words.add(other_word)
                    queue.append(other_word)
        
        # 只保留元素数>1的组
        if len(current_group) > 1:
            groups.append(current_group)
    
    return groups

关键改动说明 #

1. 用集合跟踪已分组词：grouped_words集合替代遍历组列表检查，将查询时间从O(n)降到O(1)，大幅提升效率。
2. 广度优先匹配逻辑：新增队列queue，每次取出组内的词，用它去匹配所有未分组词，只要相似度达标就加入组并继续用该词匹配，确保所有直接/间接相关的词都能被纳入同一组。
3. 提前过滤已分组词：在循环开头直接跳过已分组的词，避免重复处理。

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