基于Pandas按连续1及后续行分组生成标记列c的实现需求

原始数据 #

import pandas as pd

df = pd.DataFrame(
    {
        'a': [0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0],
        'b': [1, -1, -1, -1, 1, -1, 1, 1, -1, 1, 1, -1, 1, 1, -1, 1, 1]    
    }
)

需求 #

生成新列c，规则如下：

1. 对列a按连续1的序列以及该序列结束后的第一行进行分组；
2. 每个分组内检查列b，找到分组中第一个值为1的行；
3. 该行及之后的分组内行在c中标记为1，之前的标记为0；若分组内b没有值为1的行，则整个分组的c全为0。

期望输出 #

a  b  c
0   0  1  0
1   0 -1  0
2   1 -1  0
3   1 -1  0
4   1  1  1
5   1 -1  1
6   0  1  0
7   0  1  0
8   0 -1  0
9   1  1  1
10  1  1  1
11  1 -1  1
12  0  1  0
13  0  1  0
14  1 -1  0
15  1 -1  1
16  0  1  0

当前尝试代码（未完成） #

g = df.loc[::-1, 'a'].eq(0).cumsum()
x = df.groupby(g).filter(lambda x: x.b.iloc[0] == 1)

完整解决方案 #

步骤1：构建正确的分组键 #

你的初始思路已经接近正确，调整分组键生成逻辑，确保每个连续1序列加上后续第一个0行被划分为同一组：

# 反向遍历列a，遇到0则累加生成分组标识
g = df['a'][::-1].eq(0).cumsum()
# 反转分组标识，保证顺序与原数据一致
g = g[::-1]

步骤2：分组处理生成列c #

定义分组处理函数，对每个组执行标记逻辑：

def process_group(group):
    # 找到分组内b值为1的第一个行索引
    first_one_pos = group[group['b'] == 1].index.min()
    if pd.isna(first_one_pos):
        # 分组内无1，全标记0
        group['c'] = 0
    else:
        # 从第一个1的位置开始，后续行标记1，之前标记0
        group['c'] = (group.index >= first_one_pos).astype(int)
    return group

# 分组应用处理函数，关闭group_keys避免额外列
df = df.groupby(g, group_keys=False).apply(process_group)

验证结果 #

执行上述代码后，输出的df与期望结果完全匹配。

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