使用Pandas Groupby实现连续1分组及后续最多2个-1提取

原始数据 #

首先定义目标DataFrame：

import pandas as pd
df = pd.DataFrame({'a': [10, 20, 1, 55, 66, 333, 444, 1, 2, 10], 'b': [1,1, 1, -1, -1, -1, -1, 1, 1, -1]})

需求说明 #

需要筛选出以下行：

• 所有b列值为连续1的行
• 每个连续1序列结束后，后续最多2个b列值为-1的行（如果后续-1不足2个，则保留全部现有-1行）

期望输出结果为两个片段：

a  b
0   10  1
1   20  1
2    1  1
3   55 -1
4   66 -1

     a  b
7    1  1
8    2  1
9   10 -1

已尝试的无效代码 #

用户曾尝试以下分组方式，但未达到预期效果：

df.groupby(df.b.diff().cumsum().eq(1))
df.groupby(df['b'].ne(df['b'].shift()).cumsum())

解决方案 #

通过标记连续序列组，再对每个组后续的-1行进行计数筛选，具体步骤如下：

1. 标记连续序列的分组ID：通过b列与前一行的差异生成分组，区分不同的连续序列

# 生成连续序列的分组ID
df['group'] = df['b'].ne(df['b'].shift()).cumsum()

2. 标记每个分组所属的"1序列组"：找到所有b=1的分组，将后续的-1分组关联到前一个1分组

# 提取所有b=1的分组ID，创建映射：后续-1分组对应前一个1分组
one_groups = df[df['b'] == 1]['group'].unique()
# 为每个分组标记所属的基准1组
df['base_group'] = df['group'].apply(lambda x: max([g for g in one_groups if g <= x], default=None))

3. 对每个基准1组内的-1行进行计数筛选：保留1行全部内容，以及后续最多2个-1行

# 对每个基准1组内的行，计算属于该组的行的序号（1的行序号为0，后续-1依次递增）
df['seq_num'] = df.groupby('base_group').cumcount()
# 筛选条件：seq_num <=2（0是1的行，1和2是前两个-1行），且base_group不为空
result = df[(df['seq_num'] <= 2) & (df['base_group'].notna())].drop(['group', 'base_group', 'seq_num'], axis=1)

4. 拆分展示结果片段：按原始连续序列拆分输出

# 拆分显示结果片段
split_groups = result['b'].ne(result['b'].shift()).cumsum()
for _, g in result.groupby(split_groups):
    if not g.empty:
        print(g)
        print()

运行后输出结果与预期完全一致。

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