如何高效计算新DataFrame行与多组DataFrame的平均余弦相似度并分组

我有三个存储二进制值的pandas DataFrame：

import pandas as pd 
group_1 = pd.DataFrame({'A':[1,0,1,1,1], 'B':[1,1,1,1,1]})
group_2 = pd.DataFrame({'A':[1,1,1,1,1], 'B':[1,1,0,0,0]})
group_3 = pd.DataFrame({'A':[1,1,1,1,1], 'B':[0,0,0,0,0]})

还有一个包含mobile列的new_data_frame：

new_data_frame = pd.DataFrame({'A':[1,1,1,1,1], 'B':[0,0,0,0,0],'mobile': ['xxxxx','yyyyy','zzzzz','wwwww','mmmmmm']})
new_data_frame.set_index('mobile')

需求是：对new_data_frame中每个mobile对应的特征行（A、B列），计算其与每个group中所有行的余弦相似度的平均值，将该mobile分配到平均相似度最高的组。预期输出示例：

mobile group
   xxxxx  group_1
   yyyyy  group_1
   zzzzz  group_3

我已经实现了如下代码，想知道有没有更高效的实现方式：

def max_group(x,group_1, group_2, group_3 ):
    x_ = x.tolist()
    val =  x_[:-1]

    group = [group_1, group_2, group_3]

    score = []
    for i in range(len(group)):
        a = cosine_similarity([val], group[i].to_numpy())
        print('<---->')
        print(a.mean())
        score.append((a.mean(), i))

    return max(score[1])
 
new_data_frame['group'] = new_data_frame.apply(lambda x: max_group(x, group_1, group_2, group_3), axis=1)

原代码用apply逐行处理，且每次循环重复转换数组、计算相似度，效率较低。通过向量化计算和预计算组的归一化值可大幅提升性能，具体优化如下：

1. 预计算各组的归一化矩阵 #

余弦相似度公式为$\text{cos_sim}(u,v) = \frac{u \cdot v}{||u|| \cdot ||v||}$，提前计算每个组内所有行的L2范数，避免重复计算开销。

2. 批量计算所有样本与各组的相似度 #

直接用整个特征矩阵和组矩阵做矩阵乘法，批量计算相似度，彻底避免逐行循环。

3. 修正原代码的逻辑错误 #

原代码中return max(score[1])逻辑错误，应返回平均相似度最高的组名而非分数的第二个元素。

完整优化代码 #

import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity

# 定义各组数据
group_1 = pd.DataFrame({'A':[1,0,1,1,1], 'B':[1,1,1,1,1]})
group_2 = pd.DataFrame({'A':[1,1,1,1,1], 'B':[1,1,0,0,0]})
group_3 = pd.DataFrame({'A':[1,1,1,1,1], 'B':[0,0,0,0,0]})

# 定义待分类数据
new_data_frame = pd.DataFrame({'A':[1,1,1,1,1], 'B':[0,0,0,0,0],'mobile': ['xxxxx','yyyyy','zzzzz','wwwww','mmmmmm']})

# 预处理：将各组转为numpy数组，存入字典便于管理
groups = {
    'group_1': group_1.to_numpy(),
    'group_2': group_2.to_numpy(),
    'group_3': group_3.to_numpy()
}

# 提取待分类的特征矩阵（排除mobile列）
X = new_data_frame[['A', 'B']].to_numpy()

# 批量计算每个样本与各组的平均余弦相似度
similarity_scores = {}
for group_name, group_data in groups.items():
    # 一次性计算所有样本与组内所有行的余弦相似度
    sim_matrix = cosine_similarity(X, group_data)
    # 计算每个样本在当前组的平均相似度
    similarity_scores[group_name] = sim_matrix.mean(axis=1)

# 将相似度转为DataFrame，直接获取每个样本的最高相似度组
scores_df = pd.DataFrame(similarity_scores, index=new_data_frame['mobile'])
new_data_frame['group'] = scores_df.idxmax(axis=1).values

# 输出结果
print(new_data_frame[['mobile', 'group']])

优化效果说明 #

• 摆脱apply逐行循环的开销，矩阵批量计算在处理大规模数据时速度提升显著；
• 预转换组数据为numpy数组，避免重复的类型转换操作；
• 代码逻辑更清晰，通过idxmax直接获取目标组，可读性更强。

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