python - pandas 有条件地按组移动或向前填充另一列

Question

我有这些数据:

# df
ID     Col1 
A      -4   
A      -3
A      -2
A      -1
A       1
A       2
A    9992   
B      -4
B      -3   
B      -2
B      -1   
B       9   
B      13
B      21
B      38   
B    9992   

我想创建一个新列 Col_NEW，并基于 Col1 按 ID 进行前向填充或滚动计算。

最终结果:

# df
ID     Col1     Col_NEW
A      -4            -4
A      -3            -3
A      -2            -2
A      -1            -1
A       1             0
A       2             2
A    9992             3
B      -4            -4
B      -3            -3
B      -2            -2
B      -1            -1 
B       9             0
B      13            10
B      21            14
B      38            22
B    9992            39

逻辑如下:

1. 如果Col为-4、-3、-2、-1 code> ，则 Col_NEW 是相同的。
2. 否则，Col_NEW 将是 Col1 前向填充并加 1。

我尝试过这个，但是有没有办法直接创建Col_NEW？

df["Col_NEW"] = df.groupby('ID')['Col1'].transform(lambda x: (x.shift(1) + 1))

有什么想法吗？

Answer 1

像这样，没有任何apply或任何其他迭代:

retain_mask = df['Col1'].isin([-4,-3,-2,-1])
df["Col_NEW"] =df.groupby('ID')['Col1'].shift()+1
df.loc[retain_mask,"Col_NEW"] = df.loc[retain_mask,'Col1']

输出:

    ID      Col1    Col_NEW
--  ----  ------  ---------
 0  A         -4         -4
 1  A         -3         -3
 2  A         -2         -2
 3  A         -1         -1
 4  A          1          0
 5  A          2          2
 6  A       9992          3
 7  B         -4         -4
 8  B         -3         -3
 9  B         -2         -2
10  B         -1         -1
11  B          9          0
12  B         13         10
13  B         21         14
14  B         38         22
15  B       9992         39

使用where

同样的想法，稍微简洁一点:

retain_mask = df['Col1'].isin([-4,-3,-2,-1])
df["Col_NEW"] = df['Col1'].where(retain_mask, df.groupby('ID')['Col1'].shift()+1)