如何在Pandas中实现指定规则的累计求和生成Month列？

原始数据 #

用户的DataFrame Start Date 列数据示例：

Start Date
0
NA
NA
0
NA
NA
0
NA
0
NA
0
NA
NA
...

该列的数值统计：

df['Start Date'].value_counts(dropna=False)
# 输出：
# NaN     2145
# 0.0      130
# Name: Start Date, dtype: int64

错误尝试 #

用户尝试的代码无法得到预期结果：

df['Month'] = df['Start Date'].where(df['Start Date'] == 0, df['Start Date'].shift() + 1)

预期输出 #

期望生成的Month列如下：

Month
0
1
2
0
1
2
0
1
0
1
0
1
2
...

解决方案 #

核心思路是按Start Date中的0进行分组，在每个分组内从0开始累计计数，具体代码如下：

# 生成分组标识：每遇到一个0就开启一个新分组
df['group'] = (df['Start Date'] == 0).cumsum()
# 在每个分组内从0开始累计计数，得到目标Month列
df['Month'] = df.groupby('group').cumcount()
# 若不需要保留group列，可执行删除
# df.drop('group', axis=1, inplace=True)

原理说明 #

• (df['Start Date'] == 0).cumsum()：将Start Date等于0的位置标记为1，其余为0，累加后会为每个以0开头的连续区块生成唯一的组ID。
• groupby('group').cumcount()：在每个组内从0开始逐行计数，正好匹配需求中每个0之后依次递增1的规则。

之前的代码失效原因：使用shift()+1时，遇到连续NaN会因上一行是NaN导致结果仍为NaN，且无法实现跨分组的计数重置，无法满足每个0开头区块独立计数的要求。

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