基于状态变更获取各ID组的state=1时间窗口

核心思路 #

放弃低效的全局循环，结合Pandas分组优化+状态机逻辑，快速处理每个ID的时序序列，精准提取state=1的时间窗口，同时覆盖所有规则要求。

完整代码实现 #

import pandas as pd

def process_id_state_windows(group_df):
    # 确保数据按时间升序排列（时序数据必须保证顺序正确）
    group_df = group_df.sort_values('dt', ascending=True).reset_index(drop=True)
    
    # 规则3：无state=1的ID直接跳过
    if (group_df['state'] == 1).sum() == 0:
        return None
    
    # 规则4：无state=0的ID，直接生成唯一窗口
    if (group_df['state'] == 0).sum() == 0:
        min_dt = group_df['dt'].min()
        return pd.Series({
            'id': group_df['id'].iloc[0],
            'state1_start_dt_1': min_dt,
            'state1_end_dt_1': '2022-52'
        })
    
    # 状态机逻辑：识别state=1的连续区间
    windows = []
    current_state = None
    start_dt = None
    
    for idx, row in group_df.iterrows():
        state = row['state']
        dt = row['dt']
        
        if state == 1 and current_state != 1:
            # 进入state=1，记录起始时间
            start_dt = dt
            current_state = 1
        elif state == 0 and current_state == 1:
            # 离开state=1，结束时间为上一个时间点（当前dt是state=0的起始）
            end_dt = group_df.loc[idx-1, 'dt']
            windows.append((start_dt, end_dt))
            current_state = 0
    
    # 规则1：最后停留在state=1，补充结束时间为2022-52
    if current_state == 1:
        windows.append((start_dt, '2022-52'))
    
    # 整理为多列格式
    result = {'id': group_df['id'].iloc[0]}
    for i, (start, end) in enumerate(windows, 1):
        result[f'state1_start_dt_{i}'] = start
        result[f'state1_end_dt_{i}'] = end
    
    # 补充无对应窗口的空值
    if len(windows) < 2:
        result['state1_start_dt_2'] = None
        result['state1_end_dt_2'] = None
    
    return pd.Series(result)

# 1. 读取CSV数据（替换为你的文件路径）
df = pd.read_csv('your_data.csv')

# 2. 按ID分组处理并合并结果
result_df = df.groupby('id', group_keys=False).apply(process_id_state_windows).dropna(subset=['state1_start_dt_1'])

# 3. 输出结果到新CSV
result_df.to_csv('state_windows_result.csv', index=False)

代码说明与优化点 #

1. 时序正确性保障：每个分组先按dt排序，确保状态变化的顺序符合时间逻辑。
2. 全规则覆盖：
   • 自动过滤无state=1的ID
   • 无state=0的ID直接生成唯一窗口
   • 结尾停留在state=1的自动补2022-52
   • 多窗口自动生成带编号的列，无对应窗口的列设为NULL
3. 性能优化：用Pandas内置的groupby.apply替代手动循环ID，分组操作经过底层优化，速度远快于全局遍历；如果数据量极大，可改用矢量化状态切换识别进一步提速：

   def process_id_vectorized(group_df):
       group_df = group_df.sort_values('dt', ascending=True).reset_index(drop=True)
       
       if (group_df['state'] == 1).sum() == 0:
           return None
       if (group_df['state'] == 0).sum() == 0:
           return pd.Series({
               'id': group_df['id'].iloc[0],
               'state1_start_dt_1': group_df['dt'].min(),
               'state1_end_dt_1': '2022-52'
           })
       
       # 标记状态变化点
       group_df['state_change'] = group_df['state'] != group_df['state'].shift()
       change_points = group_df[group_df['state_change'] | (group_df.index == 0)].copy()
       
       # 处理结尾为state=1的情况
       if change_points['state'].iloc[-1] == 1:
           change_points = pd.concat([
               change_points,
               pd.DataFrame({'id': [group_df['id'].iloc[0]], 'state': [0], 'dt': ['2022-52']})
           ], ignore_index=True)
       
       # 提取state=1的起始与对应结束时间
       start_indices = change_points[change_points['state'] == 1].index
       end_indices = start_indices + 1
       windows = list(zip(change_points.loc[start_indices, 'dt'], change_points.loc[end_indices, 'dt']))
       
       # 整理结果
       result = {'id': group_df['id'].iloc[0]}
       for i, (start, end) in enumerate(windows, 1):
           result[f'state1_start_dt_{i}'] = start
           result[f'state1_end_dt_{i}'] = end
       
       if len(windows) < 2:
           result['state1_start_dt_2'] = None
           result['state1_end_dt_2'] = None
       
       return pd.Series(result)
   

   矢量化版本避免行遍历，适合超大规模数据集。

测试示例 #

假设输入CSV数据如下：

id,state,dt
A,1,2022-01
A,1,2022-02
A,0,2022-03
A,1,2022-04
B,1,2022-01
B,1,2022-02
C,0,2022-01
C,0,2022-02
D,1,2022-01
D,0,2022-02
D,0,2022-03

运行代码后输出的state_windows_result.csv内容为：

id,state1_start_dt_1,state1_end_dt_1,state1_start_dt_2,state1_end_dt_2
A,2022-01,2022-02,2022-04,2022-52
B,2022-01,2022-52,,
D,2022-01,2022-01,,

（注：CSV中空值对应规则要求的NULL）

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