首页 > 解决方案 > Pandas 在多列上滞后并设置迭代次数

问题描述

我有一个如下的数据框:

df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})

我想应用 pandas shift 函数将每一列移动 4 次,并为每个班次创建一个新行:

col1 col1.lag0 col1.lag1 col1.lag2 col1.lag3 col2 col2.lag0 col2.lag1 col2.lag2 col2.lag3
1    0         0         0         0         3    0         0         0         0
2    1         0         0         0         4    3         0         0         0
0    2         1         0         0         0    4         3         0         0
0    0         2         1         0         0    0         4         3         0
0    0         0         2         1         0    0         0         4         3

我已经尝试了一些类似 shift 的解决方案d['col1'].shift().fillna(0),但是,我不确定如何迭代解决方案,也不知道如何确保将正确数量的行添加到数据框中。

标签: pythonpandas

解决方案


首先,我将给定的 DataFrame 扩展了正确的零行数。然后迭代列和移位量以创建所需的列。

import pandas as pd

df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})
n_shifts = 4
zero_rows = pd.DataFrame(index=pd.RangeIndex(n_shift_rows), columns=df.columns).fillna(0)
df = df.append(zero_rows).reset_index(drop=True)

for col in df.columns:
    for shift_amount in range(1, n_shifts+1):
        df[f"{col}.lag{shift_amount}"] = df[col].shift(shift_amount)
df.fillna(0).astype(int)

正如 Ben.T 所指出的,可以避免外部循环,因为shift可以立即将其应用于整个 DataFrame。循环的另一种选择是

shifts = df
for shift_amount in range(1, n_shifts+1):
    columns = df.columns + ".lag" + str(shift_amount)
    shift = pd.DataFrame(df.shift(shift_amount).values, columns=columns)
    shifts = shifts.join(shift)
shifts.fillna(0).astype(int)

推荐阅读