基于多条件为Pandas数据帧的行添加标签的优雅实现

原始数据 #

import pandas as pd
import numpy as np

data = {'ID':[1]*27,
    'column1': [15, 16, 17, 14, 13, 5, 3, 2, 1.9, 1.2, 1, 0.8, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 1, 2, 3, 4, 5, 6],
    'column2': [10, 11, 12, 13, 13.5, 14, 14.5, 15, 16, 17, 18, 19, 20, 20, 20, 20, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10]
    }
df = pd.DataFrame(data)

原始DataFrame内容：

ID  column1  column2
0    1     15.0     10.0
1    1     16.0     11.0
2    1     17.0     12.0
3    1     14.0     13.0
4    1     13.0     13.5
5    1      5.0     14.0
6    1      3.0     14.5
7    1      2.0     15.0
8    1      1.9     16.0
9    1      1.2     17.0
10   1      1.0     18.0
11   1      0.8     19.0
12   1      0.5     20.0
13   1      0.5     20.0
14   1      0.5     20.0
15   1      0.5     20.0
16   1      0.5     20.0
17   1      0.5     19.0
18   1      0.5     18.0
19   1      0.5     17.0
20   1      0.5     16.0
21   1      1.0     15.0
22   1      2.0     14.0
23   1      3.0     13.0
24   1      4.0     12.0
25   1      5.0     11.0
26   1      6.0     10.0

标签赋值需求 #

• 从起始位置到首次出现column1<=2之前的行，标记为Pre_Start
• 当0.5 < column1 <= 2且column2 <= 19时，标记为Start
• 当column1 <= 0.5且column2 >= 19时，标记为Steady
• 当column1 <= 0.5且14 < column2 < 19时，标记为Ramp
• 当column1 > 0.5且column2 < 19时，标记为End

现有实现的局限性 #

此前通过分组过滤的方式仅能处理单一条件，无法适配多标签的复杂场景：

def filter_group(group):
    start_index = np.argmax(group['column1'].values <= 2)
    return group.iloc[start_index:]

filtered_df = df.groupby('ID', group_keys=False).apply(filter_group).reset_index(drop=True)

优雅实现方案 #

使用numpy.select可以清晰定义多条件映射，结合分组逻辑处理每个ID的首次column1<=2的位置，完整代码如下：

def add_labels(group):
    # 找到当前组内首次出现column1<=2的索引
    first_le2_idx = np.argmax(group['column1'] <= 2)
    
    # 定义所有条件和对应的标签
    conditions = [
        # Pre_Start：从开头到首次column1<=2的前一行
        group.index < group.index[first_le2_idx],
        # Start
        (group['column1'] > 0.5) & (group['column1'] <= 2) & (group['column2'] <= 19),
        # Steady
        (group['column1'] <= 0.5) & (group['column2'] >= 19),
        # Ramp
        (group['column1'] <= 0.5) & (group['column2'] > 14) & (group['column2'] < 19),
        # End
        (group['column1'] > 0.5) & (group['column2'] < 19)
    ]
    
    labels = [
        'Pre_Start',
        'Start',
        'Steady',
        'Ramp',
        'End'
    ]
    
    # 应用条件赋值，默认标签设为Pre_Start覆盖边界情况
    group['Label'] = np.select(conditions, labels, default='Pre_Start')
    return group

# 分组应用标签函数
result_df = df.groupby('ID', group_keys=False).apply(add_labels)

执行后得到的结果与需求示例完全一致：

ID  column1  column2      Label
0    1     15.0     10.0  Pre_Start
1    1     16.0     11.0  Pre_Start
2    1     17.0     12.0  Pre_Start
3    1     14.0     13.0  Pre_Start
4    1     13.0     13.5  Pre_Start
5    1      5.0     14.0  Pre_Start
6    1      3.0     14.5  Pre_Start
7    1      2.0     15.0      Start
8    1      1.9     16.0      Start
9    1      1.2     17.0      Start
10   1      1.0     18.0      Start
11   1      0.8     19.0      Start
12   1      0.5     20.0     Steady
13   1      0.5     20.0     Steady
14   1      0.5     20.0     Steady
15   1      0.5     20.0     Steady
16   1      0.5     20.0     Steady
17   1      0.5     19.0       Ramp
18   1      0.5     18.0       Ramp
19   1      0.5     17.0       Ramp
20   1      0.5     16.0       Ramp
21   1      1.0     15.0        End
22   1      2.0     14.0        End
23   1      3.0     13.0        End
24   1      4.0     12.0        End
25   1      5.0     11.0        End
26   1      6.0     10.0        End

方案优势 #

1. 可读性强：条件与标签一一对应，逻辑清晰易维护
2. 效率较高：基于numpy的向量化操作，比循环或逐行判断更快
3. 扩展性好：新增条件或标签只需在conditions和labels列表中添加对应项即可

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