python - 计算具有特定值的一行之后的两行

Question

我有一个包含 1 和 NaN 值的数据框，我想计算从 1 到 2 和 3 之后的两行。

import pandas as pd
df=pd.DataFrame({"b" : [1,None,None,None,None,1,None,None,None]})
print(df)

      b
0   1.0
1   NaN
2   NaN
3   NaN
4   NaN
5   1.0
6   NaN
7   NaN
8   NaN

像这样:

      b
0   1.0
1   2.0
2   3.0
3   NaN
4   NaN
5   1.0
6   2.0
7   3.0
8   NaN

我知道我可以使用 df.loc[df['b']==1] 来检索那些，但我不知道如何计算下面的两行。

Answer 1

您可以创建一个组变量，其中 b 中的每个 1 开始一个新组，然后为每个组向前填充 2 行，并执行 cumsum:

g = (df.b == 1).cumsum()
df.b.groupby(g).apply(lambda g: g.ffill(limit = 2).cumsum())

#0    1.0
#1    2.0
#2    3.0
#3    NaN
#4    NaN
#5    1.0
#6    2.0
#7    3.0
#8    NaN
#Name: b, dtype: float64