基于日期区间匹配Pandas DataFrame对应Code的高效实现

问题描述 #

我有如下结构的Pandas DataFrame：

Code    StartDate      EndDate
   A   2024-07-01   2024-08-03
   B   2024-08-06   2024-08-10
   C   2024-08-11   2024-08-31

需要遍历从2024-07-01开始的每一天，根据给定日期返回对应的Code值；若日期不在任何StartDate/EndDate区间内，则返回'Fallback_Code'。

最初的嵌套迭代实现效率极低（因DataFrame记录量大）：

DAYS = DAY_DF['Date'].tolist() # 存储所有需要匹配的日期列表
for DAY in DAYS:
    code = False
    for i,r in df.iterrows():
        if r['StartDate'] <= DAY <= r['EndDate']:
            code = r['Code']
            break
    if not code: # 如果没有匹配到任何区间
        code = 'Fallback_Code'

输入输出示例：

2024-07-03 -> 'A'
2024-08-04 -> 'Fallback_Code'
2024-08-10 -> 'B'
2024-08-11 -> 'C'

高效解决方案 #

方案一：使用pandas.merge_asof #

merge_asof是Pandas专为有序键的近似匹配设计的高效方法，适合此类区间匹配场景：

1. 转换日期类型：确保所有日期列都是datetime格式

df['StartDate'] = pd.to_datetime(df['StartDate'])
df['EndDate'] = pd.to_datetime(df['EndDate'])
DAY_DF['Date'] = pd.to_datetime(DAY_DF['Date'])

2. 排序数据集：merge_asof要求左右两边的匹配键必须排序

df_sorted = df.sort_values('StartDate')
day_df_sorted = DAY_DF.sort_values('Date')

3. 执行匹配与过滤：

# 按Date匹配最近的StartDate不大于当前日期的记录
merged = pd.merge_asof(day_df_sorted, df_sorted, left_on='Date', right_on='StartDate', direction='backward')
# 过滤掉日期超出对应EndDate的情况，赋值兜底Code
merged['Code'] = merged.apply(lambda x: x['Code'] if x['Date'] <= x['EndDate'] else 'Fallback_Code', axis=1)
# 恢复原始DAY_DF的顺序
result = merged.set_index(day_df_sorted.index).sort_index()

方案二：使用区间索引（IntervalIndex） #

利用Pandas的区间索引实现矢量级的快速匹配：

1. 转换日期并创建区间索引：

df['StartDate'] = pd.to_datetime(df['StartDate'])
df['EndDate'] = pd.to_datetime(df['EndDate'])
DAY_DF['Date'] = pd.to_datetime(DAY_DF['Date'])

# 创建包含起始和结束日期的闭区间
interval_idx = pd.IntervalIndex.from_arrays(df['StartDate'], df['EndDate'], closed='both')

2. 匹配并映射Code：

# 获取每个日期对应的区间位置，未匹配到返回-1
positions = interval_idx.get_indexer(DAY_DF['Date'])
# 映射Code，未匹配的赋值兜底值
DAY_DF['Code'] = df['Code'].iloc[positions].fillna('Fallback_Code').values

方案三：向量化布尔匹配（适合无重叠区间） #

若你的区间无重叠且数量较少，可采用批量布尔赋值的方式：

1. 转换日期类型：同上述步骤
2. 批量赋值：

# 先初始化所有Code为兜底值
DAY_DF['Code'] = 'Fallback_Code'

# 遍历每个区间，用布尔索引批量更新匹配的日期
for _, row in df.iterrows():
    mask = (DAY_DF['Date'] >= row['StartDate']) & (DAY_DF['Date'] <= row['EndDate'])
    DAY_DF.loc[mask, 'Code'] = row['Code']

此方法循环次数为区间记录数，远小于日期数时效率显著提升。

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