如何用Pandas原生方法为DataFrame按谷峰条件添加grp列？

原始DataFrame #

import numpy as np
import pandas as pd

df = pd.DataFrame(
    {'x': np.arange(1,29),
     'y': [5.69, 6.03, 6.03, 6.03, 6.03, 6.03, 6.03, 5.38, 5.21, 5.4 , 5.24,
           5.4 , 5.36, 5.47, 5.58, 5.5 , 5.61, 5.53, 5.4 , 5.51, 5.47, 5.44,
           5.39, 5.27, 5.38, 5.35, 5.32, 5.09],
     'valley': [1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
     'peak': [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,0, 0, 0, 0, 0, 0]}
)

原始输出：

x     y  valley  peak
0    1  5.69       1     0
1    2  6.03       0     0
2    3  6.03       0     1
3    4  6.03       0     0
4    5  6.03       0     0
5    6  6.03       0     0
6    7  6.03       0     0
7    8  5.38       0     0
8    9  5.21       1     0
9   10  5.40       0     0
10  11  5.24       0     0
11  12  5.40       0     0
12  13  5.36       0     0
13  14  5.47       0     0
14  15  5.58       0     0
15  16  5.50       0     0
16  17  5.61       0     1
17  18  5.53       0     0
18  19  5.40       0     0
19  20  5.51       0     0
20  21  5.47       0     0
21  22  5.44       0     0
22  23  5.39       0     0
23  24  5.27       0     0
24  25  5.38       0     0
25  26  5.35       0     0
26  27  5.32       0     0
27  28  5.09       1     0

需求说明 #

添加新列grp：

• 从valley列为1的行开始，到peak列为1的行，grp值为"A"
• 从peak列为1的行开始，到valley列为1的行，grp值为"B"

期望结果：

x     y  valley  peak  grp
0    1  5.69       1     0  A
1    2  6.03       0     0  A
2    3  6.03       0     1  B
3    4  6.03       0     0  B
4    5  6.03       0     0  B
5    6  6.03       0     0  B
6    7  6.03       0     0  B
7    8  5.38       0     0  B
8    9  5.21       1     0  A
9   10  5.40       0     0  A
10  11  5.24       0     0  A
11  12  5.40       0     0  A
12  13  5.36       0     0  A
13  14  5.47       0     0  A
14  15  5.58       0     0  A
15  16  5.50       0     0  A
16  17  5.61       0     1  B
17  18  5.53       0     0  B
18  19  5.40       0     0  B
19  20  5.51       0     0  B
20  21  5.47       0     0  B
21  22  5.44       0     0  B
22  23  5.39       0     0  B
23  24  5.27       0     0  B
24  25  5.38       0     0  B
25  26  5.35       0     0  B
26  27  5.32       0     0  B
27  28  5.09       1     0  A

解决方案（无apply/循环） #

利用Pandas的loc赋值和ffill()向前填充方法，完全原生实现：

# 初始化grp列为空值
df['grp'] = np.nan
# 标记valley触发点为"A"
df.loc[df['valley'] == 1, 'grp'] = 'A'
# 标记peak触发点为"B"
df.loc[df['peak'] == 1, 'grp'] = 'B'
# 向前填充，将触发点后的所有行继承当前分组值
df['grp'] = df['grp'].ffill()

原理说明 #

1. 先给grp列赋空值，作为初始状态
2. 用loc精准定位valley=1和peak=1的行，分别设置分组标识"A"和"B"
3. 使用ffill()（向前填充），让每个触发点后的所有行自动继承当前的分组值，直到遇到下一个新的触发点，完美匹配需求中的区间划分逻辑

验证结果 #

执行上述代码后，即可得到符合期望的grp列，与示例输出完全一致。

内容的提问来源于stack exchange，提问作者Sun Jar