基于起止布尔列的时序数据分组（Pandas向量化实现方案）

问题描述 #

现有如下结构的DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'group': ['ABC']*8 + ['DEF']*11,
    'maybe_start': [False, True, False, False, True, False, False, False,
                    False, False, True, False, False, False, False, False, True, False, False],
    'maybe_end': [False, False, False, False, False, False, True, False,
                  False, False, False, False, True, False, False, True, False, False, True]
})

需求规则：

• 按group列分组处理
• 当maybe_start为True，且是前一个maybe_end为True之后的首个True时，开启新的group2分组
• 该分组持续到首个maybe_end为True时结束，区间内所有行的group2值相同
• 不在区间内的行，group2值为NaN

预期输出：

group  maybe_start  maybe_end  group2
0    ABC        False      False     NaN
1    ABC         True      False     1.0
2    ABC        False      False     1.0
3    ABC        False      False     1.0
4    ABC         True      False     1.0
5    ABC        False      False     1.0
6    ABC        False       True     1.0
7    ABC        False      False     NaN
8    DEF        False      False     NaN
9    DEF        False      False     NaN
10   DEF         True      False     1.0
11   DEF        False      False     1.0
12   DEF        False       True     1.0
13   DEF        False      False     NaN
14   DEF        False      False     NaN
15   DEF        False       True     NaN
16   DEF         True      False     2.0
17   DEF        False      False     2.0
18   DEF        False       True     2.0

向量化实现方案 #

我们可以通过Pandas的分组、累积求和、移位和向前填充等向量化操作实现，避免循环提升效率：

# 按group分组处理
def process_group(group):
    # 1. 计算每个结束点的累积和，区分不同的结束区间
    end_cumsum = group['maybe_end'].cumsum()
    
    # 2. 标记有效的起始点：当前end区间内的第一个maybe_start
    valid_start = group['maybe_start'] & (
        group['maybe_start'].groupby(end_cumsum).transform(lambda x: x.cumsum() == 1)
    )
    
    # 3. 生成分组ID，每个有效起始对应一个递增的ID
    group_id = valid_start.cumsum()
    
    # 4. 向前填充分组ID，标记区间内的行
    active_id = group_id.where(group_id != 0).ffill()
    
    # 5. 遇到maybe_end后，后续行重置为NaN
    # 当当前行是maybe_end，下一行的active_id设为NaN
    active_id = active_id.mask(group['maybe_end'].shift(-1).fillna(False))
    
    # 6. 赋值给group2，转换为浮点型（匹配预期输出的NaN格式）
    group['group2'] = active_id.astype(float)
    
    return group

# 应用到整个DataFrame
df = df.groupby('group', group_keys=False).apply(process_group)

代码说明 #

1. end_cumsum：用maybe_end的累积和，把每个maybe_end之后的行划分到新的区间，确保只识别每个结束区间后的第一个maybe_start。
2. valid_start：筛选出每个结束区间内的第一个maybe_start，作为新分组的起始点。
3. group_id：对有效起始点进行累积求和，生成唯一的分组编号。
4. active_id：通过向前填充，把分组编号覆盖到起始点到结束点之间的所有行。
5. mask操作：在maybe_end的下一行重置分组编号为NaN，确保区间在结束点后终止。

这样就实现了完全向量化的处理，高效且符合需求。

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