python - 根据条件从当前行和上一行中的值创建列表

Question

我有一个包含两列 'a' 和 'b' 的数据框，其中 'b' 是 'a' 的值与前一个值 'a' 的差值

df = pd.DataFrame({'a': [10, 60, 30, 80, 10]})
df['b'] = df['a']-df['a'].shift(1)

    a   b
0   10  NaN
1   60  50.0
2   30  -30.0
3   80  50.0
4   10  -70.0

我想创建一个新的列 'c' ，其值是 'a' 的先前值和 'a' 的当前值的列表(例如，[60,30])，仅当列 'b' 为负时。否则，它必须是当前值 'a' 本身的列表。
结果输出应如下所示

    a   b       c
0   10  NaN     [10]
1   60  50.0    [60]
2   30  -30.0   [60, 30]
3   80  50.0    [80]
4   10  -70.0   [80, 10]

Answer 1

如果 b < 0，则使用列表理解来创建列表在带有移位辅助列的 numpy 数组中 s来自 Series.shift 添加者 DataFrame.assign :

arr = df.assign(s = df['a'].shift(fill_value=0))[['a','b','s']].to_numpy()
df['c'] = [[s,a] if b < 0 else [a] for  a,b,s in arr]
print (df)
    a     b             c
0  10   NaN        [10.0]
1  60  50.0        [60.0]
2  30 -30.0  [60.0, 30.0]
3  80  50.0        [80.0]
4  10 -70.0  [80.0, 10.0]  

或者是用 Series.mask 使用列表理解创建的一个元素列表:

s = pd.Series([[x] for x in df['a']], index=df.index)
#alternative
s = df['a'].apply(lambda x: [x])

df['c'] = s.mask(df['b'].lt(0), s.shift() + s)
print (df)
    a     b         c
0  10   NaN      [10]
1  60  50.0      [60]
2  30 -30.0  [60, 30]
3  80  50.0      [80]
4  10 -70.0  [80, 10]