基于Pandas按特定规则对DataFrame进行分组的技术咨询

现有如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame(
    {
        'a': [101, 9, 1, 10, 100, 103, 102, 105, 90, 110],
        'b': [1, 0, 1, 0, 0, 1, 1, 1, 0, 1],
    }
)

需要按列b进行分组，规则为：找到列b中满足**"1位于0之后"或"1是首行值"**的行，以这些行为起始点，创建从该行到数据末尾的分组，目标分组示例如下：

a  b
# -------------------------
0  101  1
1    9  0
2    1  1
3   10  0
4  100  0
5  103  1
6  102  1
7  105  1
8   90  0
9  110  1
# -------------------------
2    1  1
3   10  0
4  100  0
5  103  1
6  102  1
7  105  1
8   90  0
9  110  1
# -------------------------
5  103  1
6  102  1
7  105  1
8   90  0
9  110  1
# -------------------------
9  110  1

目前已通过如下代码识别了部分起始行，但未包含首行，且不知如何生成目标分组：

df.loc[(df.b == 1) & (df.b.shift(1) == 0), 'c'] = 'x'

执行后结果：

a  b    c
0  101  1  NaN
1    9  0  NaN
2    1  1    x
3   10  0  NaN
4  100  0  NaN
5  103  1    x
6  102  1  NaN
7  105  1  NaN
8   90  0  NaN
9  110  1    x

需要实现上述分组规则，以便后续对分组应用相关函数。

1. 识别所有符合条件的起始行索引 #

修正起始行的判断逻辑，把首行是1的情况纳入：

# 生成起始行的布尔掩码：首行是1，或者当前行是1且上一行是0
start_mask = (df['b'] == 1) & ((df.index == 0) | (df['b'].shift(1) == 0))
# 获取所有起始行的索引列表
start_indices = df[start_mask].index.tolist()

执行后start_indices结果为[0, 2, 5, 9]，完全匹配目标分组的起始位置。

2. 生成目标分组列表 #

基于起始索引，直接生成每个从起始行到末尾的子DataFrame：

groups = [df.loc[idx:] for idx in start_indices]

此时groups就是包含所有目标分组的列表，每个元素对应一个子DataFrame，可直接遍历使用。

3. 对分组应用函数示例 #

如果要对每个分组执行计算（比如求a列总和），直接遍历分组即可：

for i, group in enumerate(groups, 1):
    print(f"分组{i}的a列总和：{group['a'].sum()}")

输出：

分组1的a列总和：550
分组2的a列总和：442
分组3的a列总和：420
分组4的a列总和：110

可选：用groupby管理分组 #

若需要用Pandas原生的groupby结构操作，可给每个分组标记唯一ID：

# 创建组ID列，反向填充确保每个起始行到末尾的行属于对应组
df['group_id'] = None
for idx in reversed(start_indices):
    df.loc[idx:, 'group_id'] = idx

# 按group_id分组
grouped = df.groupby('group_id')

# 示例：计算每个分组的a列均值
print(grouped['a'].mean())

输出：

group_id
0      55.0
2      55.25
5      84.0
9     110.0
Name: a, dtype: float64

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