如何在Pandas DataFrame中按Stage条件计算diff_hours累计值

现有如下Pandas DataFrame：

diff_hours   stage    sensor
0            0        20
0            0        21
0            0        21
1            0        22
5            0        21
0            0        22
0            1        20
7            1        23
0            1        24
0            3        25
0            3        28
6            0        21
0            0        22

需要实现：在stage递增阶段（包括递增前的初始0阶段）累计diff_hours的值；当stage从非0状态降至0时，累计值重置为0，且后续0阶段的累计值保持0。预期结果如下：

acc_hours   stage    sensor
0            0        20
0            0        21
0            0        21
1            0        22
6            0        21
6            0        22
6            1        20
13           1        23
13           1        24
13           3        25
13           3        28
0            0        21
0            0        22

通过分组累加+条件修正的方式实现，具体代码如下：

import pandas as pd

# 构造原始DataFrame
data = {
    'diff_hours': [0,0,0,1,5,0,0,7,0,0,0,6,0],
    'stage': [0,0,0,0,0,0,1,1,1,3,3,0,0],
    'sensor': [20,21,21,22,21,22,20,23,24,25,28,21,22]
}
df = pd.DataFrame(data)

# 1. 标记重置点：当stage从非0变为0时
df['reset'] = (df['stage'] == 0) & (df['stage'].shift(1) != 0)

# 2. 生成分组键，每次重置开启新分组
df['group'] = df['reset'].cumsum()

# 3. 按分组计算diff_hours的累计和
df['acc_hours'] = df.groupby('group')['diff_hours'].cumsum()

# 4. 对重置后的0阶段行，将累计值修正为0
df.loc[(df['group'] > 0) & (df['stage'] == 0), 'acc_hours'] = 0

# 保留需要的列并输出
result = df[['acc_hours', 'stage', 'sensor']]
print(result)

1. 标记重置点：通过shift(1)获取前一行的stage值，判断当前行是否是从非0切换到0的情况，标记为重置点。
2. 生成分组键：对重置点进行累加，得到分组标识，确保每次重置后的数据属于新分组。
3. 分组累计求和：按分组对diff_hours做累计求和，实现同分组内的累加逻辑。
4. 修正重置后的0阶段值：对重置后的分组（即group>0）中的0阶段行，将累计值强制设为0，符合需求中的重置规则。

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