在Pandas中基于序列内满足条件的连续出现次数创建新列

给定如下Pandas时间序列数据集，需按account分组新增output列，满足特定连续值计数规则：

import pandas as pd

date0 = ["2020-01-31"] * 3
date1 = ["2020-02-28"] * 3
date2 = ["2020-03-31"] * 3

date = date0 + date1 + date2
acc = ["A", "B", "C"] * 3
val = [1, 0, 0, 1, 1, 0, 1, 1, 1]

df = pd.DataFrame({"date": date, "account": acc, "value": val})

规则说明 #

1. 当value列连续出现3次（可自定义N次）1时，对应行的output设为1，否则为0
2. 若连续达标后出现value=0，则连续计数重置；若后续继续出现value=1，则重新累计，再次达到N次时output设为1；若连续达标后持续为1，output保持1

示例验证 #

• 原数据集预期输出：

  output = [0, 0, 0, 0, 0, 0, 1, 0, 0]
  df["output"] = output
  

• 扩展数据集验证（后续时间点规则）：

  date3 = ["2020-04-30"] * 3
  val = [0, 1, 1]
  acc = ["A", "B", "C"]
  output = [0, 1, 0]
  
  df2 = pd.DataFrame({"date": date3, "account": acc, "value": val, "output": output})
  df3 = pd.concat([df, df2])
  

以下方案基于Pandas内置分组操作实现，无自定义迭代/函数，适配大型数据集：

import pandas as pd

# 定义连续达标阈值
N = 3

# 第一步：确保数据按账户+时间排序（关键前提）
df = df.sort_values(['account', 'date']).reset_index(drop=True)

# 第二步：标记连续1的中断点（value=0时生成新分组ID）
mask = df['value'] == 0
df['reset_group'] = mask.groupby(df['account']).cumsum()

# 第三步：计算每个连续组内的1的计数，value=0时计数重置为0
df['run_length'] = df.groupby(['account', 'reset_group']).cumcount() + 1
df['run_length'] = df['run_length'] * df['value']  # 清零value=0的计数

# 第四步：生成output列
df['output'] = (df['run_length'] >= N).astype(int)

方案解释 #

1. 排序：必须先按account和date排序，保证时间序列的顺序正确性
2. 中断点标记：通过reset_group列将每个账户的连续1序列划分为独立分组，遇到value=0时分组ID递增，实现计数重置
3. 连续长度计算：用cumcount统计每个分组内的位置，再与value相乘，将非1的行计数清零，得到真实的连续1长度
4. 生成结果：直接通过阈值比较生成output，完全向量化操作，性能优异

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