如何结合groupby与重置条件，用cumcount/cumsum实现DataFrame计数重置？

现有如下Pandas DataFrame，包含ID和Tag两列：

import pandas as pd
df = pd.DataFrame({'ID': ['A', 'A', 'A', 'A', 'B', 'B', 'B'], 'Tag': ['', 'X', '', '', 'X', '','X']}  )  

需要生成Counts列，规则为：从0开始依次递增，当ID发生变更或Tag取值为X时，计数重新从0开始。

解决方法 #

核心思路是先确定所有需要重置计数的位置，再按这些位置划分分组，最后在每个分组内从0开始顺序计数。

代码实现 #

# 标记需要重置计数的行：ID变更 或 当前Tag为X
reset_flag = (df['ID'] != df['ID'].shift()) | (df['Tag'] == 'X')
# 生成分组标识，每遇到一次重置就开启新分组
group_key = reset_flag.cumsum()
# 每个分组内从0开始计数
df['Counts'] = df.groupby(group_key).cumcount()

运行结果 #

执行代码后，DataFrame变为：

步骤说明 #

1. 重置标记：用shift()对比当前行与上一行的ID是否不同，同时判断当前行Tag是否为X，得到所有需要重置计数的行标记。
2. 分组标识：对重置标记做累加，每出现一次重置标记，分组标识加1，这样每个重置位置都会成为新分组的起点。
3. 分组计数：通过groupby按分组标识分组，再用cumcount()在每个分组内从0开始依次计数，完美匹配需求。

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