python - 如何加速这个Python函数？

Question

我希望提出此类问题是可以的。

我有一个 get_lags 函数，它接受一个数据帧，并且对于每一列，将列移动列表 n_lags 中的每个 n。因此，如果 n_lags = [1, 2]，该函数会将每列移动 1 个位置和 2 个位置，从而以这种方式创建新的滞后列。

def get_lags (df, n_lags):
    data =df.copy()
    data_with_lags = pd.DataFrame()
    for column in data.columns:
        for i in range(n_lags[0], n_lags[-1]+1):
            new_column_name = str(column) + '_Lag' + str(i)
            data_with_lags[new_column_name] = data[column].shift(-i) 
    data_with_lags.fillna(method = 'ffill', limit = max(n_lags), inplace = True)
    return data_with_lags

所以，如果:

df.columns
ColumnA
ColumnB

然后，get_lags(df, [1 , 2]).columns 将是:

ColumnA_Lag1
ColumnA_Lag2
ColumnB_Lag1
ColumnB_Lag2

问题:使用具有大约 100,000 行和 20,000 列的数据框，这需要很长时间才能运行。在一台 16 GB RAM、Core i7 Windows 机器上，有一次我等待代码运行 15 分钟才停止它。我是否可以调整此功能以使其更快？

Answer 1

您需要shift + concat。这是简洁的版本 -

def get_lags(df, n_lags):
    return pd.concat(
       [df] + [df.shift(i).add_suffix('_Lag{}'.format(i)) for i in n_lags],
       axis=1
    )

这是一个更内存友好的版本，使用 for 循环 -

def get_lags(df, n_lags):
    df_list = [df]
    for i in n_lags:
        v = df.shift(i)
        v.columns = v.columns + '_Lag{}'.format(i)    
        df_list.append(v)

    return pd.concat(df_list, axis=1)