Python Pandas：创建带条件重置的累计变量列

我明白你要做的事儿——就是要追踪连续戒咖啡的天数，遇到forgot的时候不打断累计 streak，喝了咖啡就直接重置为0对吧？我之前也帮别人处理过类似的累计追踪问题，给你两种解决方案，一种是容易理解的自定义函数法，另一种是更适合大数据量的矢量化方法，你可以按需选择：

先准备示例数据 #

首先咱们先模拟你提到的场景，创建一个包含日期和每日状态的DataFrame：

import pandas as pd

# 构造示例数据
data = {
    'date': pd.date_range(start='2024-01-01', periods=7),
    'status': ['no_coffee', 'no_coffee', 'forgot', 'had_coffee', 'no_coffee', 'forgot', 'no_coffee']
}
df = pd.DataFrame(data)

方法一：自定义分组函数（易理解） #

这种方法逻辑清晰，适合刚接触Pandas的朋友，核心思路是把每次喝了咖啡的时刻作为分组边界，在每个分组内逐个计算连续天数：

# 第一步：创建分组标识，每次遇到had_coffee就开启一个新分组
df['group'] = (df['status'] == 'had_coffee').cumsum()

# 第二步：定义函数计算每个分组内的连续戒咖啡天数
def calculate_streak(group):
    current_streak = 0
    streak_list = []
    for status in group['status']:
        if status == 'had_coffee':
            current_streak = 0  # 喝了咖啡，重置为0
        elif status == 'no_coffee':
            current_streak += 1  # 没喝咖啡，天数+1
        # 遇到forgot，保持当前天数不变，不用额外处理
        streak_list.append(current_streak)
    group['coffee_free_streak'] = streak_list
    return group

# 应用函数到每个分组，最后删除临时分组列
df = df.groupby('group').apply(calculate_streak).drop('group', axis=1)

方法二：矢量化运算（高效） #

如果你的数据量很大，用矢量化方法会比循环快很多，核心是利用cumsum做分组累计，巧妙处理forgot的情况：

# 标记需要重置的点（喝了咖啡的行）
reset_points = df['status'] == 'had_coffee'
# 创建分组ID：每次重置后开启新分组
df['group'] = reset_points.cumsum()
# 临时列：no_coffee标记为1，其他（forgot/had_coffee）标记为0
df['temp'] = (df['status'] == 'no_coffee').astype(int)

# 分组内累计求和，自动处理forgot的保持逻辑
df['coffee_free_streak'] = df.groupby('group')['temp'].cumsum()
# 把喝了咖啡的行强制设为0
df.loc[reset_points, 'coffee_free_streak'] = 0

# 清理临时列
df = df.drop(['group', 'temp'], axis=1)

最终结果 #

两种方法都会得到如下结果：

你可以根据自己的实际数据调整状态的判断条件，比如如果还有其他状态，只需要在逻辑里补充对应的处理规则就行~

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