Pandas中带重置功能的条件累积计数实现方案优化

你的需求是对布尔列做累积计数，遇到0时重置为0，原循环方法在大型DataFrame中性能较差，这里提供基于Pandas矢量化操作的高效实现方案，完全避免Python循环，性能提升显著。

实现代码 #

import pandas as pd

# 示例DataFrame
df = pd.DataFrame({'bool': [1,0,1,1,1,1,1,0,1,1,1,0,0,1,1,1]})

# 步骤1：生成分组键，连续非0值归为同一组，遇0则新建分组
groups = df['bool'].eq(0).cumsum()

# 步骤2：组内累积计数并处理0值
df['counter'] = df.groupby(groups).cumcount().add(1) * df['bool']

# 查看结果
print(df)

代码解释 #

1. 分组键生成：df['bool'].eq(0).cumsum() 会将所有等于0的位置标记为True，累加后每遇到一个0，分组编号就加1，这样连续的非0值会被分到同一个组，每个0或连续0会成为独立分组。
2. 组内计数与重置：groupby(groups).cumcount().add(1) 对每个组内元素从1开始计数，再乘以原bool列，自动将原列为0的位置的计数重置为0，完美符合需求。

更紧凑的写法 #

可以把两步合并为一行代码，效果完全一致：

df['counter'] = df.groupby(df['bool'].eq(0).cumsum()).cumcount().add(1).where(df['bool'] != 0, 0)

性能说明 #

对于百万级甚至千万级行的DataFrame，这种矢量化方法的速度是Python循环的几十到上百倍，因为底层基于numpy的C语言实现，避免了逐元素的Python级循环开销。

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