如何基于True掩码及后续N行对Pandas DataFrame分组？

原始数据 #

import pandas as pd
df = pd.DataFrame(
    {
        'a': [False, True, False, True, False, True, False, True, True, False, False],
    }
)

需求说明 #

需要将df.a == True的行及其后续两行划分为一组，且跳过已被之前分组覆盖的行：

• 第一个组从第1行（第一个True）开始，包含行1、2、3
• 第二个组找不在第一个组中的下一个True（第5行），包含行5、6、7
• 第三个组找不在前两个组中的下一个True（第8行），包含行8、9、10

预期分组输出：

a
1    True
2   False
3    True

        a
5    True
6   False
7    True

      a
8    True
9   False
10  False

尝试的无效代码 #

N = 2
mask = ((df.a.eq(True))
        .cummax().cumsum()
        .between(1, N+1)
        )

out = df[mask]

解决方案代码 #

import pandas as pd

df = pd.DataFrame(
    {
        'a': [False, True, False, True, False, True, False, True, True, False, False],
    }
)

# 1. 定位所有a为True的行索引
true_indices = df[df['a']].index.tolist()
selected_starts = []
last_group_end = -1

# 2. 筛选出未被之前分组覆盖的起始True索引
for idx in true_indices:
    if idx > last_group_end:
        selected_starts.append(idx)
        last_group_end = idx + 2  # 每组覆盖当前行+后续2行

# 3. 拆分并输出每个分组
for start in selected_starts:
    group = df.loc[start:start+2]
    print(group)
    print()

代码逻辑说明 #

1. 先找出所有a为True的行索引，再筛选出未被之前分组覆盖的起始索引（当前True索引需大于上一个分组的结束位置）
2. 对每个筛选后的起始索引，直接截取当前行及后续两行作为分组
3. 循环打印每个分组，得到预期的输出格式

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