如何加速Python嵌套循环？Pandas大数据集高效替代方案咨询

需求与数据集说明 #

• 数据集一（date_status）：包含两列，唯一标识ID，以及每个ID对应的嵌套DataFrame（含Date列、默认值为1的Status列，1代表无事件）
• 数据集二（event_data）：包含ID、Begin Date、End Date列，记录各ID的事件时间区间
• 核心需求：遍历每个ID，检查其嵌套DataFrame中的日期是否落在对应ID的任意事件区间内，若是则将Status改为0（代表有事件）

数据集示例 #

数据集一（date_status）

数据集二（event_data）

现有问题 #

当前使用三层嵌套循环实现需求，但面对2万+ID、18万+事件、5年每日数据的规模，运行时间超过12小时，完全无法满足效率要求。尝试过where()、mask()、apply()等方法，但未找到正确传递索引的方式，急需向量化的高效解决方案。

低效代码示例 #

def check_ID(id, id_event):
    cs = pd.DataFrame(date_status.loc[id:id,].drop(columns=['status']))
    event_id = pd.DataFrame(event_data.loc[id_event:id_event,].drop(columns=['Begin Date','End Date','Consecutive_Days']))
    return cs.equals(event_id)

def is_event_true(id_index, id_event_index, date_index):
    date_val = date_status.loc[id_index,'status'].loc[date_index,'Date']
    begin = event_data.loc[id_event_index,'Begin Date']
    end = event_data.loc[id_event_index,'End Date']
    return (date_val >= begin) and (date_val < end)

for i in date_status.index:
    for x in event_data.index:
        if check_ID(i,x):
            for y in date_status.loc[i,'status'].index:
                if is_event_true(i,x,y):
                    date_status.loc[i,'status'].loc[y,'Status'] = 0

高效解决方案：扁平化数据 + 向量化区间匹配 #

嵌套DataFrame是性能瓶颈的核心原因，先将数据集一扁平化，再通过向量化操作完成区间匹配，步骤如下：

1. 扁平化数据集一 #

把每个ID对应的嵌套DataFrame展开，和外层ID合并成一个大的DataFrame：

import pandas as pd
import numpy as np

# 展开嵌套的Status列
flattened_ds = pd.concat(
    [pd.Series({'ID': row['ID']}).repeat(len(row['status'])), row['status']]
    for _, row in date_status.iterrows()
).reset_index(drop=True)

# 确保ID列和日期列类型正确
flattened_ds['ID'] = flattened_ds['ID'].astype(int)
flattened_ds['Date'] = pd.to_datetime(flattened_ds['Date'])

2. 标准化数据集二的日期格式 #

event_data['Begin Date'] = pd.to_datetime(event_data['Begin Date'])
event_data['End Date'] = pd.to_datetime(event_data['End Date'])

3. 向量化区间匹配 #

方案A：交叉合并+布尔过滤（适合中等规模数据）

# 按ID合并两个数据集
merged = flattened_ds.merge(event_data, on='ID', how='left')

# 标记日期是否在事件区间内
merged['in_event'] = (merged['Date'] >= merged['Begin Date']) & (merged['Date'] < merged['End Date'])

# 按ID和Date分组，只要有一个事件区间包含该日期，就将Status设为0
flattened_ds['Status'] = merged.groupby(['ID', 'Date'])['in_event'].transform(lambda x: 0 if x.any() else 1)

方案B：merge_asof优化（适合超大规模事件数据）

# 先将扁平化数据集按ID和Date排序
flattened_ds_sorted = flattened_ds.sort_values(['ID', 'Date'])
# 事件数据集按ID和Begin Date排序
event_data_sorted = event_data.sort_values(['ID', 'Begin Date'])

# 按ID匹配，找到每个Date对应的最近的Begin Date不晚于Date的事件
merged = pd.merge_asof(
    flattened_ds_sorted,
    event_data_sorted,
    on='Date',
    by='ID',
    direction='backward'
)

# 检查Date是否在匹配到的事件区间内，更新Status
flattened_ds_sorted['Status'] = np.where(
    (merged['Date'] >= merged['Begin Date']) & (merged['Date'] < merged['End Date']),
    0,
    1
)

# 若需要恢复原嵌套结构，可重新分组：
date_status_updated = flattened_ds_sorted.groupby('ID').apply(
    lambda x: x[['Date', 'Status']].reset_index(drop=True)
).reset_index().rename(columns={0: 'status'})

性能优势 #

扁平化后所有操作都是向量化的，避免了Python层面的循环，对于2万+ID、5年每日数据（约365*5=1825条/ID，总约3650万条）和18万+事件的规模，运行时间可控制在数分钟内，远优于原嵌套循环方案。

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