如何用Pandas按首列'B'分组DataFrame？（含起始'B'，不含结束'B'）

需求描述 #

需要对包含数千条记录的DataFrame进行分组：

• 分组从首列值为'B'的行开始
• 到下一个首列值为'B'的行结束（结束的'B'行不包含在当前分组内）

输入示例数据 #

B,0,15000.000000,716.881652,-0.065916 
K,0,-33,1,4030
K,1,-16,2,4028
K,2,-18,12,4036
K,3,-14,-3,4054
P,0,-452,4089,329
P,1,-428,4082,427
P,2,-382,4078,518
P,3,-363,4052,545
P,4,-347,4064,508
K,4,-2,17,4048
K,5,-18,12,4048
P,5,-373,4068,409
B,1,16000.000000,715.443420,-0.067204
K,6,-16,-2,4054
P,6,-433,4082,390
K,7,-16,4,4036
P,7,-457,4104,406
B,2,17000.000000,716.930297,-0.084369

期望输出 #

• df1:

B,0,15000.000000,716.881652,-0.065916 
K,0,-33,1,4030
K,1,-16,2,4028
K,2,-18,12,4036
K,3,-14,-3,4054
P,0,-452,4089,329
P,1,-428,4082,427
P,2,-382,4078,518
P,3,-363,4052,545
P,4,-347,4064,508
K,4,-2,17,4048
K,5,-18,12,4048
P,5,-373,4068,409

• df2:

B,1,16000.000000,715.443420,-0.067204
K,6,-16,-2,4054
P,6,-433,4082,390
K,7,-16,4,4036
P,7,-457,4104,406

实现步骤与代码 #

核心思路 #

1. 定位所有首列为'B'的行索引，作为分组的起始标记
2. 为每行分配对应的分组ID，确保每个分组覆盖从当前'B'行到下一个'B'行前的所有内容
3. 过滤掉最后一个不完整的分组（最后一个'B'行之后无后续'B'行，无法形成有效分组）
4. 按分组ID拆分得到目标DataFrame

完整代码 #

import pandas as pd
from io import StringIO

# 模拟输入数据（实际使用时替换为pd.read_csv("你的文件路径.csv")）
data = """B,0,15000.000000,716.881652,-0.065916 
K,0,-33,1,4030
K,1,-16,2,4028
K,2,-18,12,4036
K,3,-14,-3,4054
P,0,-452,4089,329
P,1,-428,4082,427
P,2,-382,4078,518
P,3,-363,4052,545
P,4,-347,4064,508
K,4,-2,17,4048
K,5,-18,12,4048
P,5,-373,4068,409
B,1,16000.000000,715.443420,-0.067204
K,6,-16,-2,4054
P,6,-433,4082,390
K,7,-16,4,4036
P,7,-457,4104,406
B,2,17000.000000,716.930297,-0.084369"""

# 加载数据到DataFrame
df = pd.read_csv(StringIO(data), header=None)

# 获取所有首列为'B'的行索引
b_indices = df[df[0] == 'B'].index.tolist()

# 为每行生成分组ID
group_ids = []
current_group = 0
for idx in df.index:
    # 遇到'B'行时更新分组ID（最后一个'B'行不开启新分组）
    if idx in b_indices and idx != b_indices[-1]:
        current_group = b_indices.index(idx)
    group_ids.append(current_group)

df['group'] = group_ids

# 过滤掉最后一个无效分组，拆分得到目标DataFrame
valid_groups = df[df['group'] < len(b_indices)-1]
grouped_dfs = [group.drop('group', axis=1) for _, group in valid_groups.groupby('group')]

# 提取结果
df1 = grouped_dfs[0]
df2 = grouped_dfs[1]

# 可选：验证输出
# print("df1:\n", df1.to_csv(header=False, index=False))
# print("\ndf2:\n", df2.to_csv(header=False, index=False))

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