如何在Pandas中基于code列生成连续1分组的Period列？

原始代码 #

import pandas as pd

# 示例数据
data = { 'name': ['ray', 'ray', 'ray', 'ray', 'ray', 'ray'], 
            'code': [1, 0, 1, 1, 0, 1] }

# 创建DataFrame
df = pd.DataFrame(data)

# 初始化'Period'列
df['Period'] = 0

# 按逻辑计算'Period'
current_period = 0 
for i in range(len(df)): 
    if df.loc[i, 'code'] == 1: 
        current_period += 1 
    else: 
        current_period = 0 
    df.loc[i, 'Period'] = current_period

# 输出带'Period'列的DataFrame
print(df)

需求说明 #

需要根据code列的值生成Period列，具体规则：

• 第一行code=1，Period为1；
• 第二行code=0，Period为0；
• 第3、4行连续code=1，两行的Period均为2；
• 第5行code=0，Period为0；
• 第6行code=1，Period为3，以此类推。

核心是把连续的code=1行划分为一组，为每组赋予递增的编号，code=0的行Period固定为0。

问题分析 #

当前代码是逐行累加计数，导致连续的code=1行生成1、2这类逐行递增的数值，而非同一组统一的编号，不符合需求。

解决方案 #

通过标记连续code=1的分组，再为每组分配编号的方式实现：

import pandas as pd

# 示例数据
data = { 'name': ['ray', 'ray', 'ray', 'ray', 'ray', 'ray'], 
            'code': [1, 0, 1, 1, 0, 1] }
df = pd.DataFrame(data)

# 标记连续code=1的分组：当code从0变为1时，分组编号+1
df['group'] = (df['code'] == 1) & (df['code'].shift(1) != 1)
df['group'] = df['group'].cumsum()

# 生成Period列：code=1时取分组编号，否则为0
df['Period'] = df.apply(lambda x: x['group'] if x['code'] == 1 else 0, axis=1)

# 删除临时分组列
df = df.drop('group', axis=1)

print(df)

运行后输出：

name  code  Period
0   ray     1       1
1   ray     0       0
2   ray     1       2
3   ray     1       2
4   ray     0       0
5   ray     1       3

代码解释 #

1. 标记分组：先筛选code=1的行，再判断该行是否是连续1组的起始行（上一行code≠1），最后用cumsum()累加得到分组编号；
2. 生成Period列：通过apply函数判断，code=1时取对应分组编号，code=0时设为0；
3. 清理临时列：删除用于分组的临时group列，得到最终结果。

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