Pandas无循环实现遇0重置的counter计数列

核心思路 #

用向量化运算替代显式循环，先给每段连续计数区间打分组标记，再做组内计数即可：

• 识别所有counter列值为0的重置点
• 对重置点做累计求和生成分组ID：每遇到一个0，分组ID加1，同一段连续区间内的所有行共享同一个分组ID
• 按分组ID做组内从0开始的顺序编号，直接得到目标计数值

实现代码 #

先构造测试数据集：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    "date": pd.date_range("2017-08-01", "2017-08-02", freq='H'),
    "counter": [0, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 0, np.nan, np.nan, np.nan, np.nan, np.nan, 0, np.nan, np.nan, 0, 0, np.nan, np.nan, np.nan, np.nan, 0, np.nan, np.nan]
})

核心计算逻辑仅需2行：

# 生成区间分组标记
df['group_tag'] = df['counter'].eq(0).cumsum()
# 组内从0开始顺序计数，覆盖原counter列
df['counter'] = df.groupby('group_tag').cumcount()

如果不需要保留中间列，计算完成后直接删除即可：

df = df.drop(columns=['group_tag'])

结果校验 #

执行代码后得到的结果和预期完全一致：

date  counter
0  2017-08-01 00:00:00        0
1  2017-08-01 01:00:00        1
2  2017-08-01 02:00:00        2
3  2017-08-01 03:00:00        3
4  2017-08-01 04:00:00        4
5  2017-08-01 05:00:00        5
6  2017-08-01 06:00:00        6
7  2017-08-01 07:00:00        0
8  2017-08-01 08:00:00        1
9  2017-08-01 09:00:00        2
10 2017-08-01 10:00:00        3
11 2017-08-01 11:00:00        4
12 2017-08-01 12:00:00        5
13 2017-08-01 13:00:00        0
14 2017-08-01 14:00:00        1
15 2017-08-01 15:00:00        2
16 2017-08-01 16:00:00        0
17 2017-08-01 17:00:00        0
18 2017-08-01 18:00:00        1
19 2017-08-01 19:00:00        2
20 2017-08-01 20:00:00        3
21 2017-08-01 21:00:00        4
22 2017-08-01 22:00:00        0
23 2017-08-01 23:00:00        1
24 2017-08-02 00:00:00        2

原理解释 #

• eq(0)会将列中值为0的位置标记为布尔值True（运算时等价于1），其余位置标记为False（运算时等价于0）
• 接cumsum()做累计求和时，每碰到一个0值，累计和就加1，两个相邻0值之间的所有行累计和保持不变，自动完成区间划分，不需要手动循环判断
• cumcount()是Pandas内置的组内顺序计数方法，默认从0开始编号，全程走Pandas底层向量化实现，性能远高于Python层手写循环，百万行级数据也能毫秒级返回结果。

内容的提问来源于stack exchange，提问作者john-mueller