Python中如何在flag列值变为1时拆分pandas DataFrame

需求规则 #

当flag列的取值变化为1时，在该位置对DataFrame切分，切分后丢弃所有flag=1的行，仅保留切分得到的非1值连续数据段。

测试数据 #

import pandas as pd
df = pd.DataFrame({
    'A':[1,20,40,45,56,1,20,40,45,56],
    'flag':[3,2,4,1,1,3,3,1,1,1]
})

原始DataFrame打印结果：

A  flag
0   1     3
1  20     2
2  40     4
3  45     1
4  56     1
5   1     3
6  20     3
7  40     1
8  45     1
9  56     1

期望输出 #

• group_1：

A  flag
0   1     3
1  20     2
2  40     4

• group_2：

A  flag
0   1     3
1  20     3

实现代码 #

核心逻辑是先识别切分边界生成分组ID，过滤无效行后按分组ID拆分即可：

# 识别切分边界：当前行flag为1、上一行flag不为1时，判定为一个分组的结束点，累加分组ID
df['group_tag'] = (df['flag'].eq(1) & df['flag'].shift().ne(1)).cumsum()
# 过滤掉flag=1的行，按分组标签拆分后重置索引
split_groups = [sub_df.reset_index(drop=True) for _, sub_df in df[df['flag'] != 1].groupby('group_tag')]

# 按顺序取出分组
group_1 = split_groups[0]
group_2 = split_groups[1]

执行后打印group_1、group_2即可得到完全匹配期望的结果。

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