python - 如何加快大数据集中的两个嵌套 for 循环

Question

我正在处理一个 600.000 x 128 的数据集，按以下形式排列，并带有层次索引（见左图）：

                    Var1     Var2     ...                                     Var1_mean  Var2_mean     ...          
    ID     Date                                               ID     Date          
    1      2017-12  1        0.1                              1      2017-12  2          0.3
    1      2018-01  2        0.3                              1      2018-01  2          0.3
    1      2018-02  3        0.5            -->               1      2018-02  2          0.3
    2      2018-01  2        0.2                              2      2018-01  3          0.15
    2      2018-02  4        0.1                              2      2018-02  3          0.15

现在对于所有 55.000 个客户（= 唯一 ID），我想用每个客户的平均值替换所选变量中每个客户的值，存储在新变量中（变量名称 + '_mean'，见右图）。

我写了一个函数来做到这一点，但它需要 4 个多小时，即使在尝试并行化任务之后也是如此。我试过了：

    identifiers = set(df_raw_sort.index.get_level_values(0)) # unique IDs

    def avg_per_customer(column): 
       df_raw_sort.loc[:, column + '_mean'] = 0 # Create new col
       for unique in identifiers: 
          meanvalue = np.mean(df_raw_sort[column].loc[(unique)])
          df_raw_sort.loc[(unique), column + '_mean'] = meanvalue

    Parallel(n_jobs=2, verbose=10)(delayed(avg_per_customer)(col) for col in transform)
    # transforms selected columns (= transform)

我能做些什么来加快这件事？

提前非常感谢。

标签： pythonperformanceloopsfor-loopparallel-processing

Answer 1

你能试试这个

df_raw_sort[variable_name_mean] = df_raw_sort[variable_name].groupby(df_raw_sort['ID']).transform('mean')

从这里参考答案