Pandas：当year列重置为1时统计分组编号

我有一个大型DataFrame，需要在year列重置为1时对分组进行编号统计。尝试了以下代码但未得到预期结果：

df["count"] = df.groupby((df["year"]>=df["year"].shift(1)).cumsum()).cumcount()+1

示例DataFrame：

year month   alpha
1      0     1.24036
1      1     1.17023
1      0     1.24036
1      1     1.17023
1      0     1.24036
1      1     1.17023
2      0     1.24036
2      1     1.17023
2      0     1.24036
2      1     1.17023
2      0     1.24036
2      1     1.17023
1      0     1.24036
1      1     1.17023
1      0     1.24036
1      1     1.17023
1      0     1.24036
1      1     1.17023
2      0     1.24036
2      1     1.17023
2      0     1.24036
2      1     1.17023
2      0     1.24036
2      1     1.17023

预期输出：

year month   alpha     count
1      0     1.24036     1
1      1     1.17023     1
1      0     1.24036     1    
1      1     1.17023     1
1      0     1.24036     1
1      1     1.17023     1
2      0     1.24036     1
2      1     1.17023     1
2      0     1.24036     1
2      1     1.17023     1
2      0     1.24036     1
2      1     1.17023     1
1      0     1.24036     2
1      1     1.17023     2
1      0     1.24036     2 
1      1     1.17023     2
1      0     1.24036     2
1      1     1.17023     2
2      0     1.24036     2
2      1     1.17023     2
2      0     1.24036     2
2      1     1.17023     2
2      0     1.24036     2
2      1     1.17023     2

你原来的代码存在两个问题：

1. 分组条件错误：df["year"]>=df["year"].shift(1)仅判断当前年份不小于上一年份，当年份从2变回1时，该条件不成立，无法触发新分组的创建。
2. 使用cumcount()+1会在每个分组内逐行生成递增序号，而你需要的是整个分组共用同一个编号。

我们需要先识别出新分组的起始点（即year从非1值变为1的位置），再通过累加起始点标记来生成分组编号：

# 标记新分组的起始行：当当前year为1且上一行year不为1时，标记为True
group_start = (df["year"] == 1) & (df["year"].shift() != 1)
# 对起始标记累加，得到每个行对应的分组编号
df["count"] = group_start.cumsum()

这段代码会直接为每个完整的分组（包含连续的1和2序列）分配统一的编号，完全符合你的预期输出。

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