arrays - 基于前 n 行在 groupby() 中创建新列的更短方法

Question

我有以下代码，用于已排序的 Pandas 数据框，按一列分组，并创建两列新列:一列根据组中的前 4 行和当前行，另一列基于组中的 future 行。

data_test = {'nr':[1,1,1,1,1,6,6,6,6,6,6,6],
            'val':[11,12,13,14,15,61,62,63,64,65,66,67]}
df_test = pd.DataFrame (data_test, columns = ['nr','val'])

print (df_test) 

因此有以下框架:

   nr  val
0    1   11
1    1   12
2    1   13
3    1   14
4    1   15
5    6   61
6    6   62
7    6   63
8    6   64
9    6   65
10   6   66
11   6   67

现在我必须遵循按“nr”分组的代码，并为每一行构建一列，其中包含组中“val”的前 4 个值和当前值。同样，构建一个额外的列，每行包含组中 'val' 的 future 值。

df_test['past4'] = df_test.groupby(['nr'])['val'].transform(lambda x: x.shift(4).fillna(0))
df_test['past3'] = df_test.groupby(['nr'])['val'].transform(lambda x: x.shift(3).fillna(0))
df_test['past2'] = df_test.groupby(['nr'])['val'].transform(lambda x: x.shift(2).fillna(0))
df_test['past1'] = df_test.groupby(['nr'])['val'].transform(lambda x: x.shift(1).fillna(0))
df_test['future'] = df_test.groupby(['nr'])['val'].transform(lambda x: x.shift(-1).fillna(0))
df_test['amounts'] = df_test[['past4', 'past3','past2','past1','val']].values.tolist()
df_test.drop(columns = ['past4', 'past3', 'past2', 'past1'], inplace = True)
df_test

    nr  val future  amounts
0   1   11  12  [0, 0, 0, 0, 11]
1   1   12  13  [0, 0, 0, 11, 12]
2   1   13  14  [0, 0, 11, 12, 13]
3   1   14  15  [0, 11, 12, 13, 14]
4   1   15  0   [11, 12, 13, 14, 15]
5   6   61  62  [0, 0, 0, 0, 61]
6   6   62  63  [0, 0, 0, 61, 62]
7   6   63  64  [0, 0, 61, 62, 63]
8   6   64  65  [0, 61, 62, 63, 64]
9   6   65  66  [61, 62, 63, 64, 65]
10  6   66  67  [62, 63, 64, 65, 66]
11  6   67  0   [63, 64, 65, 66, 67]

我确信我应该能够更容易地构建一个名为“数量”的列表列，可能是单行的。我怎样才能做到这一点？

Answer 1

使用自定义函数创建嵌套列表，例如:

def f(x):
    #list comprehension with shift by 4,3,2,1,0
    L = [x['val'].shift(i).fillna(0) for i in range(4, -1, -1)]
    #shifting to another column
    x['future'] = x['val'].shift(-1).fillna(0).astype(int)
    #column filled by lists
    x['amounts'] = pd.Series(np.array(L).astype(int).T.tolist(), index=x.index)
    return (x)

df_test = df_test.groupby(['nr']).apply(f)
print (df_test)
    nr  val  future               amounts
0    1   11      12      [0, 0, 0, 0, 11]
1    1   12      13     [0, 0, 0, 11, 12]
2    1   13      14    [0, 0, 11, 12, 13]
3    1   14      15   [0, 11, 12, 13, 14]
4    1   15       0  [11, 12, 13, 14, 15]
5    6   61      62      [0, 0, 0, 0, 61]
6    6   62      63     [0, 0, 0, 61, 62]
7    6   63      64    [0, 0, 61, 62, 63]
8    6   64      65   [0, 61, 62, 63, 64]
9    6   65      66  [61, 62, 63, 64, 65]
10   6   66      67  [62, 63, 64, 65, 66]
11   6   67       0  [63, 64, 65, 66, 67]