如何对二维数组进行向量化运算计算余弦相似度？

问题原因 #

scipy.spatial.distance.cosine仅支持计算一对1维向量的余弦距离，直接传入2维数组（整列数据）会触发维度不匹配报错；而apply逐行处理本质是Python循环，在大型DataFrame上效率极低。

最优解决方案：Numpy向量化实现 #

直接基于余弦相似度公式手动实现向量化计算，完全规避循环，这是速度最快的方案：
余弦相似度公式：$\text{sim} = \frac{\vec{A} \cdot \vec{B}}{||\vec{A}|| \times ||\vec{B}||}$

代码示例：

import numpy as np
import pandas as pd

# 生成示例数据（固定随机种子保证可复现）
np.random.seed(42)
df = pd.DataFrame(np.random.randn(20).reshape(5,4), columns=["ref_x", "ref_y", "alt_x", "alt_y"])

# 提取两组向量的numpy数组
ref_vectors = df[["ref_x", "ref_y"]].values
alt_vectors = df[["alt_x", "alt_y"]].values

# 计算每行对应向量的点积
dot_products = (ref_vectors * alt_vectors).sum(axis=1)

# 计算每组向量的模长
norm_ref = np.linalg.norm(ref_vectors, axis=1)
norm_alt = np.linalg.norm(alt_vectors, axis=1)

# 计算余弦相似度，添加极小值避免除以0的情况
epsilon = 1e-10
df["sim"] = dot_products / (norm_ref * norm_alt + epsilon)

备选方案：使用Sklearn成对相似度计算 #

如果偏好现成库函数，可借助sklearn.metrics.pairwise.cosine_similarity，它会计算两个矩阵所有行向量的两两相似度，取对角线元素即可得到对应行的向量相似度：

from sklearn.metrics.pairwise import cosine_similarity

df["sim"] = cosine_similarity(ref_vectors, alt_vectors).diagonal()

注：该方案代码更简洁，但效率略低于手动Numpy实现，且依赖Sklearn库。

性能对比 #

• apply方法：Python循环开销大，百万级行数据处理速度极慢；
• Numpy向量化：底层C实现，速度比apply提升100~1000倍；
• Sklearn方法：速度介于两者之间，适合追求代码简洁的场景。

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