利用numpy.vectorize提升计算速度

同步于音尘杂记

前言

在实际项目中，对超大矩阵进行计算或者对超大的DataFrame进行计算是一个经常会出现的场景。这里先不考虑开发机本身内存等客观硬件因素，仅从设计上讨论一下不同实现方式带来的性能差异，抛砖引玉。

项目中有这样一个需求，需要根据历史销量数据计算SKU(Stock Keeping Unit)之间的相似度，或者更通俗一点说是根据历史销量数据求不同SKU之间出现的订单交集以及并集大小(注:SKU数量大概15k左右，订单数大概1000k左右)。

这里给几条示例数据，可以更直观形象地理解这个需求：

然后需要根据这些历史的orderno-sku(订单-商品)数据求解出sku的相似度矩阵。其中SKU1和SKU2之间的相似度定义为:

可以很快速地想到几种解决方案：

• 直接for loops；
• for loops稍微改进采用列表生成器；
• 采用多进程并行计算；
• 采用numpy.vectorize

1. for loops计算相似度矩阵

@timer
def gen_corr_matrix_for_loops(order_df):
    """
    for loops计算相似度矩阵
    """

    df = order_df.groupby(['sku']).agg({
   'orderno': lambda x: set(x)}).reset_index()
    del order_df
    gc.collect()
    l = len(df)
    sku_series = df.sku.astype(str)
    corr_matrix_arr = np.ones((l, l))

    tbar = trange(l)
    tbar.set_description("compute corr matrix")
    for i in tbar:
        for j in range(i + 1, l):
            corr_matrix_arr[j, i] = corr_matrix_arr[i, j] = len(df.iloc[i, 1] & df.iloc[j, 1]) / len(
                df.iloc[i, 1] | df.iloc[j, 1])
    corr_matrix_df = pd.DataFrame