如何基于DataFrame列的连续条件创建子DataFrame？

问题描述 #

给定如下DataFrame，需要基于B列中**连续的"1"**生成独立的子DataFrame。预期得到三个子DataFrame，分别对应原DataFrame的第1-2行、第5行、第7-10行。

示例数据 #

import pandas as pd

level_1 = ['A', 'B']
data = [['a1', 1], ['a2', 1],['a3', 0], ['a4', 0],['a1', 1], ['a5', 0],['a6', 1], ['a7', 1],['a8', 1], ['a9', 1]]
df = pd.DataFrame(data, columns=level_1)

解决方案 #

核心思路是先给连续的1段落打上分组标识，再按分组拆分出子DataFrame：

1. 生成分组标识：利用ne(1)标记非1的行，再通过cumsum()为连续的1段落分配唯一分组ID；
2. 筛选并分组：筛选出B=1的行，按分组ID拆分得到子DataFrame列表。

完整代码：

# 生成连续1的分组标识
df['group'] = df['B'].ne(1).cumsum()

# 拆分得到子DataFrame列表（移除辅助分组列）
sub_dfs = [group_df.drop(columns='group') for _, group_df in df[df['B'] == 1].groupby('group')]

# 查看结果
for idx, sub_df in enumerate(sub_dfs, 1):
    print(f"子DataFrame {idx}:")
    print(sub_df)
    print("-"*25)

输出结果 #

子DataFrame 1:
    A  B
0  a1  1
1  a2  1
-------------------------
子DataFrame 2:
    A  B
4  a1  1
-------------------------
子DataFrame 3:
    A  B
6  a6  1
7  a7  1
8  a8  1
9  a9  1
-------------------------

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