首页 > 解决方案 > 获取映射值的余弦距离的有效(不是 DataFrame.apply)方法

问题描述

这是我生成的一些数据:

import numpy as np
import pandas as pd
import scipy
import scipy.spatial

df = pd.DataFrame(
    {
        "item_1": np.random.randint(low=0, high=10, size=1000),
        "item_2": np.random.randint(low=0, high=10, size=1000),
    }
)
embeddings = {item_id: np.random.randn(100) for item_id in range(0, 10)}


def get_distance(item_1, item_2):
    arr1 = embeddings[item_1]
    arr2 = embeddings[item_2]
    return scipy.spatial.distance.cosine(arr1, arr2)

我想申请get_distance每一行。我可以:

df.apply(lambda row: get_distance(row["item_1"], row["item_2"]), axis=1)

但这对于大型数据集来说会非常慢。

有没有办法在不使用的情况下计算与每一行对应的嵌入的余弦相似度DataFrame.apply

标签: pythonpandasnumpyscipycosine-similarity

解决方案


直接使用向量化numpy操作要快得多:

item_1_embedded = np.array([embeddings[x]for x in df.item_1])
item_2_embedded = np.array([embeddings[x]for x in df.item_2])
cos_dist = 1 - np.sum(item_1_embedded*item_2_embedded, axis=1)/(np.linalg.norm(item_1_embedded, axis=1)*np.linalg.norm(item_2_embedded, axis=1))

(这个版本771 µs平均在我的电脑上运行,与37.4 ms. DataFrame.apply,这使得纯 numpy 版本快约 50 倍)。


推荐阅读