基于连续12期状态列生成Flag的高效代码优化方案咨询

示例数据 #

规则说明 #

对于每个唯一ID，标记某一行的Flag时，需要检查该行后续12个月的status列：如果这段时间内存在status=1的记录，该行Flag输出1，否则输出0。

原低效实现代码 #

原代码通过嵌套循环处理，但面对百万级数据时性能极差：

df = df1[['ID','Period','status']]
df = df.sort_values(by=['ID','Period']).reset_index(drop=True)
df['flag']=0
default=pd.DataFrame()

j=df['ID'].unique()
k=pd.DataFrame(j,columns=['ID'])
for a,i in k.iterrows():
      temp = df[df['ID']==i['ID']]
      temp = temp.reset_index(drop=True)
 
      for x,y in temp.iterrows():
            if temp.iloc[x+1:x+13,2].sum() >= 1:
                 temp.loc[x,'flag']=1
            else:
                 temp.loc[x,'flag']=0
      default = default.append(temp)

高效实现方案 #

利用Pandas的滚动窗口+反向计算可以大幅提升性能，完全避免嵌套循环：

代码实现 #

import pandas as pd

# 转换日期类型，确保排序逻辑正确
df['Period'] = pd.to_datetime(df['Period'])
# 按ID和日期排序，保证时间顺序
df = df.sort_values(by=['ID', 'Period']).reset_index(drop=True)

# 核心逻辑：分组后反向滚动窗口求和，判断是否存在status=1
df['flag'] = df.groupby('ID')['status'].apply(
    lambda s: s[::-1].rolling(window=12, min_periods=1).sum()[::-1] > 0
).astype(int)

性能优势 #

• 完全基于Pandas向量化运算和C级优化，避免Python层面的循环开销
• 百万级数据处理时间从分钟级压缩到秒级
• 代码简洁，逻辑清晰易维护

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